From d22a53f17e7cf50fce6ec8ab5f5bc44eedef4ba2 Mon Sep 17 00:00:00 2001 From: jetie <1405758738@qq.com> Date: Fri, 12 Jun 2026 21:05:01 +0800 Subject: [PATCH] feat(primitive): add barrier and queue primitives Add distributed barrier and FIFO queue primitives following the same pattern as lock and election. Barrier provides synchronization: create holds a barrier that blocks other callers until released. Supports lease-based auto-release and watch-based blocking semantics (create_with_wait). Queue provides FIFO operations: enqueue adds items with sequential keys, dequeue pops the front item atomically. Peek reads without removing. Both primitives emit watch events on state changes and integrate with the Raft state machine, lease system, and gRPC API layer. --- build.rs | 2 + proto/barrier.proto | 60 +++++++ proto/queue.proto | 58 +++++++ src/api/barrier.rs | 241 ++++++++++++++++++++++++++ src/api/mod.rs | 4 + src/api/queue.rs | 146 ++++++++++++++++ src/barrier/mod.rs | 325 +++++++++++++++++++++++++++++++++++ src/lib.rs | 2 + src/main.rs | 30 +++- src/queue/mod.rs | 241 ++++++++++++++++++++++++++ src/raft/mod.rs | 25 +++ src/raft/state_machine.rs | 347 +++++++++++++++++++++++++++++++++++++- 12 files changed, 1478 insertions(+), 3 deletions(-) create mode 100644 proto/barrier.proto create mode 100644 proto/queue.proto create mode 100644 src/api/barrier.rs create mode 100644 src/api/queue.rs create mode 100644 src/barrier/mod.rs create mode 100644 src/queue/mod.rs diff --git a/build.rs b/build.rs index 01ec4e6..69c39b4 100644 --- a/build.rs +++ b/build.rs @@ -11,6 +11,8 @@ fn main() -> Result<(), Box> { "proto/maintenance.proto", "proto/lock.proto", "proto/election.proto", + "proto/barrier.proto", + "proto/queue.proto", ], &["proto/"], )?; diff --git a/proto/barrier.proto b/proto/barrier.proto new file mode 100644 index 0000000..74e02f6 --- /dev/null +++ b/proto/barrier.proto @@ -0,0 +1,60 @@ +syntax = "proto3"; + +package aether; + +import "common.proto"; + +// AetherBarrier provides distributed barrier operations. +// A barrier is a synchronization primitive: while the barrier exists, +// other callers that try to enter it will block until the barrier +// is released or its lease expires. +service AetherBarrier { + // Create creates and holds a barrier with the given name. + // Other callers trying to enter will block until this barrier + // is released or its lease expires. + // Returns a barrier key that can be used to release or query. + rpc Create(BarrierCreateRequest) returns (BarrierCreateResponse); + // Release releases the barrier, unblocking all waiters. + rpc Release(BarrierReleaseRequest) returns (BarrierReleaseResponse); + // Query returns whether the barrier is currently held. + rpc Query(BarrierQueryRequest) returns (BarrierQueryResponse); +} + +message BarrierCreateRequest { + // name is the barrier name. + bytes name = 1; + // lease_id is the lease to associate with the barrier. + // If 0, the barrier will not be associated with a lease + // and must be explicitly released. + int64 lease_id = 2; +} + +message BarrierCreateResponse { + ResponseHeader header = 1; + // key is the key created for the barrier. + bytes key = 2; +} + +message BarrierReleaseRequest { + // name is the barrier name to release. + bytes name = 1; +} + +message BarrierReleaseResponse { + ResponseHeader header = 1; +} + +message BarrierQueryRequest { + // name is the barrier name to query. + bytes name = 1; +} + +message BarrierQueryResponse { + ResponseHeader header = 1; + // held is true if the barrier is currently held. + bool held = 2; + // key is the barrier key if held, empty otherwise. + bytes key = 3; + // lease_id is the lease associated with the barrier, 0 if none. + int64 lease_id = 4; +} diff --git a/proto/queue.proto b/proto/queue.proto new file mode 100644 index 0000000..55a82c0 --- /dev/null +++ b/proto/queue.proto @@ -0,0 +1,58 @@ +syntax = "proto3"; + +package aether; + +import "common.proto"; + +// AetherQueue provides distributed FIFO queue operations. +// Items are stored with sequential keys under a named queue prefix. +// Enqueue appends an item; Dequeue pops the front item atomically. +service AetherQueue { + // Enqueue adds an item to the named queue. + rpc Enqueue(QueueEnqueueRequest) returns (QueueEnqueueResponse); + // Dequeue removes and returns the front item from the named queue. + // Returns NOT_FOUND if the queue is empty. + rpc Dequeue(QueueDequeueRequest) returns (QueueDequeueResponse); + // Peek returns the front item without removing it. + // Returns NOT_FOUND if the queue is empty. + rpc Peek(QueuePeekRequest) returns (QueuePeekResponse); +} + +message QueueEnqueueRequest { + // name is the queue name. + bytes name = 1; + // value is the item to enqueue. + bytes value = 2; +} + +message QueueEnqueueResponse { + ResponseHeader header = 1; + // key is the key assigned to the enqueued item. + bytes key = 2; +} + +message QueueDequeueRequest { + // name is the queue name. + bytes name = 1; +} + +message QueueDequeueResponse { + ResponseHeader header = 1; + // key is the key of the dequeued item. + bytes key = 2; + // value is the value of the dequeued item. + bytes value = 3; +} + +message QueuePeekRequest { + // name is the queue name. + bytes name = 1; +} + +message QueuePeekResponse { + ResponseHeader header = 1; + // key is the key of the front item. + bytes key = 2; + // value is the value of the front item. + bytes value = 3; +} diff --git a/src/api/barrier.rs b/src/api/barrier.rs new file mode 100644 index 0000000..f037f78 --- /dev/null +++ b/src/api/barrier.rs @@ -0,0 +1,241 @@ +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; + +use tonic::{Request, Response, Status}; + +use crate::barrier::{BarrierManager, barrier_key, validate_barrier_name}; +use crate::proto::aether_barrier_server::AetherBarrier; +use crate::proto::{ + BarrierCreateRequest, BarrierCreateResponse, BarrierQueryRequest, BarrierQueryResponse, + BarrierReleaseRequest, BarrierReleaseResponse, ResponseHeader, +}; +use crate::raft::{self, RaftHandle, WatchEventType, require_leader}; +use crate::watch::WatchManager; + +/// Maximum time to wait for barrier acquisition (30 seconds). +const BARRIER_TIMEOUT: Duration = Duration::from_secs(30); + +pub struct BarrierService { + raft: Arc, + node_id: u64, + barrier_manager: Arc>, + watch_manager: Arc, +} + +impl BarrierService { + pub fn new( + raft: Arc, + node_id: u64, + barrier_manager: Arc>, + watch_manager: Arc, + ) -> Self { + Self { + raft, + node_id, + barrier_manager, + watch_manager, + } + } + + fn header(&self) -> ResponseHeader { + ResponseHeader { + cluster_id: 0, + member_id: self.node_id, + revision: 0, + raft_term: self.raft.term(), + } + } + + async fn propose(&self, request: raft::RaftRequest) -> Result { + require_leader(self.raft.as_ref(), self.node_id)?; + self.raft + .propose(request) + .await + .map_err(|e| Status::internal(format!("raft write failed: {e}"))) + } + + /// Try to create a barrier, blocking if it's already held. + /// Creates a watch BEFORE proposing to avoid missing release events. + async fn create_with_wait(&self, name: Vec, lease_id: i64) -> Result, Status> { + // Create a watch on the barrier key BEFORE proposing + // This ensures we don't miss any Delete events + let key = barrier_key(&name); + let (watch_id, mut watch_rx) = self + .watch_manager + .create( + key, + Vec::new(), // exact key match + vec![WatchEventType::Delete], + false, + ) + .await; + + // Helper to cancel watch + let cancel_watch = || { + let wm = self.watch_manager.clone(); + async move { + wm.cancel(watch_id, "barrier create completed".to_string()) + .await; + } + }; + + loop { + // Propose barrier create to Raft + let resp = match self + .propose(raft::RaftRequest::BarrierCreate { + name: name.clone(), + lease_id, + }) + .await + { + Ok(resp) => resp, + Err(e) => { + cancel_watch().await; + return Err(e); + } + }; + + match resp { + raft::RaftResponse::BarrierCreate { key } => { + cancel_watch().await; + return Ok(key); + } + raft::RaftResponse::BarrierAlreadyHeld { .. } => { + // Wait for the barrier key to be deleted + tracing::debug!( + barrier = %String::from_utf8_lossy(&name), + "barrier is held, waiting for release" + ); + + // Wait for the key to be deleted or timeout + let result = tokio::time::timeout(BARRIER_TIMEOUT, async { + while let Some(resp) = watch_rx.recv().await { + if resp.canceled { + break; + } + for event in &resp.events { + if event.event_type == WatchEventType::Delete { + return true; + } + } + } + false + }) + .await; + + match result { + Ok(true) => { + // Barrier was released, retry create + tracing::debug!( + barrier = %String::from_utf8_lossy(&name), + "barrier released, retrying create" + ); + continue; + } + Ok(false) => { + // Watch was canceled unexpectedly + cancel_watch().await; + return Err(Status::internal("watch canceled unexpectedly")); + } + Err(_) => { + // Timeout + cancel_watch().await; + return Err(Status::deadline_exceeded(format!( + "barrier create timed out after {} seconds: barrier is held", + BARRIER_TIMEOUT.as_secs() + ))); + } + } + } + raft::RaftResponse::Error { message } => { + cancel_watch().await; + return Err(Status::internal(message)); + } + _ => { + cancel_watch().await; + return Err(Status::internal("unexpected response type")); + } + } + } + } +} + +#[tonic::async_trait] +impl AetherBarrier for BarrierService { + async fn create( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + + validate_barrier_name(&req.name).map_err(Status::invalid_argument)?; + + if req.lease_id < 0 { + return Err(Status::invalid_argument("lease_id must be non-negative")); + } + + // Use create_with_wait to block if barrier is already held + let key = self.create_with_wait(req.name, req.lease_id).await?; + + Ok(Response::new(BarrierCreateResponse { + header: Some(self.header()), + key, + })) + } + + async fn release( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + + validate_barrier_name(&req.name).map_err(Status::invalid_argument)?; + + let resp = self + .propose(raft::RaftRequest::BarrierRelease { name: req.name }) + .await?; + + match resp { + raft::RaftResponse::BarrierRelease {} => Ok(Response::new(BarrierReleaseResponse { + header: Some(self.header()), + })), + raft::RaftResponse::Error { message } => Err(Status::internal(message)), + _ => Err(Status::internal("unexpected response type")), + } + } + + async fn query( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + + validate_barrier_name(&req.name).map_err(Status::invalid_argument)?; + + require_leader(self.raft.as_ref(), self.node_id)?; + + let mgr = self + .barrier_manager + .lock() + .map_err(|e| Status::internal(format!("barrier manager lock poisoned: {e}")))?; + + let key = mgr + .get_key(&req.name) + .map(|k| k.to_vec()) + .unwrap_or_default(); + let held = !key.is_empty(); + let lease_id = if held { + mgr.get_lease_id(&key).unwrap_or(0) + } else { + 0 + }; + + Ok(Response::new(BarrierQueryResponse { + header: Some(self.header()), + held, + key, + lease_id, + })) + } +} diff --git a/src/api/mod.rs b/src/api/mod.rs index e24725d..ca238b6 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -1,4 +1,5 @@ mod auth; +mod barrier; mod cluster; mod election; pub mod health; @@ -7,15 +8,18 @@ mod lease; mod lock; pub mod maintenance; pub mod metrics; +mod queue; mod shard; mod watch; pub use self::auth::AuthService; +pub use self::barrier::BarrierService; pub use self::cluster::ClusterService; pub use self::election::ElectionService; pub use self::kv::KvService; pub use self::lease::LeaseService; pub use self::lock::LockService; pub use self::maintenance::MaintenanceService; +pub use self::queue::QueueService; pub use self::shard::ShardService; pub use self::watch::WatchService; diff --git a/src/api/queue.rs b/src/api/queue.rs new file mode 100644 index 0000000..b484a15 --- /dev/null +++ b/src/api/queue.rs @@ -0,0 +1,146 @@ +use std::sync::Arc; + +use tonic::{Request, Response, Status}; + +use crate::proto::aether_queue_server::AetherQueue; +use crate::proto::{ + QueueDequeueRequest, QueueDequeueResponse, QueueEnqueueRequest, QueueEnqueueResponse, + QueuePeekRequest, QueuePeekResponse, ResponseHeader, +}; +use crate::queue::{queue_scan_prefix, validate_queue_name}; +use crate::raft::{self, RaftHandle, require_leader}; +use crate::storage::StorageEngine; + +/// Maximum value size for queue items (1 MB). +const MAX_QUEUE_VALUE_SIZE: usize = 1024 * 1024; + +pub struct QueueService { + raft: Arc, + node_id: u64, + storage: Arc, +} + +impl QueueService { + pub fn new( + raft: Arc, + node_id: u64, + storage: Arc, + ) -> Self { + Self { + raft, + node_id, + storage, + } + } + + fn header(&self) -> ResponseHeader { + ResponseHeader { + cluster_id: 0, + member_id: self.node_id, + revision: 0, + raft_term: self.raft.term(), + } + } + + async fn propose(&self, request: raft::RaftRequest) -> Result { + require_leader(self.raft.as_ref(), self.node_id)?; + self.raft + .propose(request) + .await + .map_err(|e| Status::internal(format!("raft write failed: {e}"))) + } +} + +#[tonic::async_trait] +impl AetherQueue for QueueService { + async fn enqueue( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + + validate_queue_name(&req.name).map_err(Status::invalid_argument)?; + + if req.value.is_empty() { + return Err(Status::invalid_argument("value must not be empty")); + } + if req.value.len() > MAX_QUEUE_VALUE_SIZE { + return Err(Status::invalid_argument(format!( + "value size {} exceeds maximum {}", + req.value.len(), + MAX_QUEUE_VALUE_SIZE + ))); + } + + let resp = self + .propose(raft::RaftRequest::QueueEnqueue { + name: req.name, + value: req.value, + }) + .await?; + + match resp { + raft::RaftResponse::QueueEnqueue { key } => Ok(Response::new(QueueEnqueueResponse { + header: Some(self.header()), + key, + })), + raft::RaftResponse::Error { message } => Err(Status::internal(message)), + _ => Err(Status::internal("unexpected response type")), + } + } + + async fn dequeue( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + + validate_queue_name(&req.name).map_err(Status::invalid_argument)?; + + let resp = self + .propose(raft::RaftRequest::QueueDequeue { name: req.name }) + .await?; + + match resp { + raft::RaftResponse::QueueDequeue { key, value } => { + Ok(Response::new(QueueDequeueResponse { + header: Some(self.header()), + key, + value, + })) + } + raft::RaftResponse::QueueDequeueEmpty {} => Err(Status::not_found("queue is empty")), + raft::RaftResponse::Error { message } => Err(Status::internal(message)), + _ => Err(Status::internal("unexpected response type")), + } + } + + async fn peek( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + + validate_queue_name(&req.name).map_err(Status::invalid_argument)?; + + require_leader(self.raft.as_ref(), self.node_id)?; + + // Read the front item directly from storage (without proposing through Raft) + let prefix = queue_scan_prefix(&req.name); + let entries = self + .storage + .scan(&prefix, 1) + .map_err(|e| Status::internal(format!("storage scan failed: {e}")))?; + + if entries.is_empty() { + return Err(Status::not_found("queue is empty")); + } + + let front = &entries[0]; + Ok(Response::new(QueuePeekResponse { + header: Some(self.header()), + key: front.key.clone(), + value: front.value.clone(), + })) + } +} diff --git a/src/barrier/mod.rs b/src/barrier/mod.rs new file mode 100644 index 0000000..cd010a4 --- /dev/null +++ b/src/barrier/mod.rs @@ -0,0 +1,325 @@ +use std::collections::{HashMap, HashSet}; + +use crate::storage::{RocksStorage, StorageEngine}; + +/// Key prefix for barrier keys in storage. +pub const BARRIER_KEY_PREFIX: &[u8] = b"_aether_barrier/"; + +/// Reserved prefix for system use. +pub const RESERVED_PREFIX: &[u8] = b"_aether_"; + +/// Maximum allowed barrier name length (1 KB). +pub const MAX_BARRIER_NAME_LEN: usize = 1024; + +/// Lease ID size in bytes (i64 = 8 bytes). +const LEASE_ID_SIZE: usize = 8; + +/// In-memory barrier manager. Owned by the state machine. +/// +/// Tracks which barriers are currently held. The actual barrier data is stored +/// in the KV store with the `_aether_barrier/` prefix. +/// +/// Value format in KV store: `[lease_id: i64 BE][name_bytes]` +pub struct BarrierManager { + /// Active barriers: barrier_name -> barrier_key + barriers: HashMap, Box<[u8]>>, + /// Reverse index: barrier_key -> barrier_name + key_to_name: HashMap, Box<[u8]>>, + /// Lease association: barrier_key -> lease_id + key_lease: HashMap, i64>, + /// Reverse index: lease_id -> set of barrier_keys + lease_keys: HashMap>>, +} + +impl Default for BarrierManager { + fn default() -> Self { + Self::new() + } +} + +impl BarrierManager { + pub fn new() -> Self { + Self { + barriers: HashMap::new(), + key_to_name: HashMap::new(), + key_lease: HashMap::new(), + lease_keys: HashMap::new(), + } + } + + /// Restore barrier state from persistent storage. + pub fn restore(&mut self, storage: &RocksStorage) -> Result<(), crate::error::StorageError> { + let entries = storage.scan(BARRIER_KEY_PREFIX, usize::MAX)?; + + self.barriers.clear(); + self.key_to_name.clear(); + self.key_lease.clear(); + self.lease_keys.clear(); + + for entry in entries { + let key: Box<[u8]> = entry.key.into(); + let value = entry.value; + + if value.len() < LEASE_ID_SIZE { + tracing::warn!( + key = %String::from_utf8_lossy(&key), + "invalid barrier value, skipping" + ); + continue; + } + + let lease_id_bytes: [u8; 8] = match value[..LEASE_ID_SIZE].try_into() { + Ok(bytes) => bytes, + Err(_) => { + tracing::warn!( + key = %String::from_utf8_lossy(&key), + "failed to convert lease_id bytes, skipping" + ); + continue; + } + }; + let lease_id = i64::from_be_bytes(lease_id_bytes); + let name: Box<[u8]> = value[LEASE_ID_SIZE..].into(); + + self.barriers.insert(name.clone(), key.clone()); + self.key_to_name.insert(key.clone(), name); + if lease_id > 0 { + self.key_lease.insert(key.clone(), lease_id); + self.lease_keys.entry(lease_id).or_default().insert(key); + } + } + + Ok(()) + } + + /// Encode barrier value with lease_id. + pub fn encode_value(name: &[u8], lease_id: i64) -> Vec { + let mut value = Vec::with_capacity(LEASE_ID_SIZE + name.len()); + value.extend_from_slice(&lease_id.to_be_bytes()); + value.extend_from_slice(name); + value + } + + /// Create a barrier. Returns the barrier key. + pub fn create(&mut self, name: Vec, key: Vec, lease_id: i64) -> Vec { + let name_box: Box<[u8]> = name.into(); + let key_box: Box<[u8]> = key.into(); + + if lease_id > 0 { + self.key_lease.insert(key_box.clone(), lease_id); + self.lease_keys + .entry(lease_id) + .or_default() + .insert(key_box.clone()); + } else if let Some(old_lease_id) = self.key_lease.remove(&key_box) + && let Some(keys) = self.lease_keys.get_mut(&old_lease_id) + { + keys.remove(&key_box); + if keys.is_empty() { + self.lease_keys.remove(&old_lease_id); + } + } + + self.key_to_name.insert(key_box.clone(), name_box.clone()); + self.barriers.insert(name_box, key_box.clone()); + + key_box.into_vec() + } + + /// Release a barrier by key. Returns true if found and released. + pub fn release(&mut self, key: &[u8]) -> bool { + if let Some(name) = self.key_to_name.remove(key) { + self.barriers.remove(&name); + if let Some(lease_id) = self.key_lease.remove(key) + && let Some(keys) = self.lease_keys.get_mut(&lease_id) + { + keys.remove(key); + if keys.is_empty() { + self.lease_keys.remove(&lease_id); + } + } + true + } else { + false + } + } + + /// Release all barriers associated with a lease. + pub fn release_by_lease(&mut self, lease_id: i64) -> Vec> { + let keys: Vec> = match self.lease_keys.remove(&lease_id) { + Some(keys) => keys.into_iter().collect(), + None => return Vec::new(), + }; + + let mut result = Vec::with_capacity(keys.len()); + for key in keys { + if let Some(name) = self.key_to_name.remove(&key) { + self.barriers.remove(&name); + } + self.key_lease.remove(&key); + result.push(key.into_vec()); + } + + result + } + + /// Check if a barrier is held. + pub fn is_held(&self, name: &[u8]) -> bool { + self.barriers.contains_key(name) + } + + /// Get the barrier key for a given name. + pub fn get_key(&self, name: &[u8]) -> Option<&[u8]> { + self.barriers.get(name).map(|k| k.as_ref()) + } + + /// Get the lease_id for a given barrier key. + pub fn get_lease_id(&self, key: &[u8]) -> Option { + self.key_lease.get(key).copied() + } + + /// Get the number of active barriers. + pub fn barrier_count(&self) -> usize { + self.barriers.len() + } + + /// Get all barrier keys associated with a lease. + pub fn get_keys_by_lease(&self, lease_id: i64) -> Vec> { + match self.lease_keys.get(&lease_id) { + Some(keys) => keys.iter().map(|k| k.to_vec()).collect(), + None => Vec::new(), + } + } +} + +/// Generate a barrier key for a given name. +pub fn barrier_key(name: &[u8]) -> Vec { + let mut key = Vec::with_capacity(BARRIER_KEY_PREFIX.len() + name.len()); + key.extend_from_slice(BARRIER_KEY_PREFIX); + key.extend_from_slice(name); + key +} + +/// Extract the barrier name from a barrier key. +pub fn barrier_name(key: &[u8]) -> Option<&[u8]> { + if key.starts_with(BARRIER_KEY_PREFIX) { + Some(&key[BARRIER_KEY_PREFIX.len()..]) + } else { + None + } +} + +/// Validate barrier name. +pub fn validate_barrier_name(name: &[u8]) -> Result<(), &'static str> { + if name.is_empty() { + return Err("barrier name must not be empty"); + } + if name.len() > MAX_BARRIER_NAME_LEN { + return Err("barrier name too long"); + } + if name.contains(&0) { + return Err("barrier name must not contain null bytes"); + } + if name.starts_with(RESERVED_PREFIX) { + return Err("barrier name must not start with reserved prefix '_aether_'"); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_barrier_key_generation() { + let name = b"my-barrier"; + let key = barrier_key(name); + assert_eq!(key, b"_aether_barrier/my-barrier"); + } + + #[test] + fn test_barrier_name_extraction() { + let key = b"_aether_barrier/my-barrier"; + let name = barrier_name(key); + assert_eq!(name, Some(b"my-barrier".as_slice())); + } + + #[test] + fn test_barrier_name_extraction_invalid() { + let key = b"not-a-barrier"; + let name = barrier_name(key); + assert_eq!(name, None); + } + + #[test] + fn test_encode_decode_value() { + let name = b"test-barrier"; + let lease_id = 42i64; + let encoded = BarrierManager::encode_value(name, lease_id); + + assert_eq!(encoded.len(), LEASE_ID_SIZE + name.len()); + assert_eq!(&encoded[..LEASE_ID_SIZE], &lease_id.to_be_bytes()); + assert_eq!(&encoded[LEASE_ID_SIZE..], name); + } + + #[test] + fn test_barrier_manager_create_release() { + let mut mgr = BarrierManager::new(); + let name = b"test-barrier".to_vec(); + let key = barrier_key(&name); + + assert!(!mgr.is_held(&name)); + assert_eq!(mgr.barrier_count(), 0); + + mgr.create(name.clone(), key.clone(), 0); + assert!(mgr.is_held(&name)); + assert_eq!(mgr.get_key(&name), Some(key.as_slice())); + assert_eq!(mgr.barrier_count(), 1); + + assert!(mgr.release(&key)); + assert!(!mgr.is_held(&name)); + assert_eq!(mgr.barrier_count(), 0); + } + + #[test] + fn test_barrier_manager_release_nonexistent() { + let mut mgr = BarrierManager::new(); + assert!(!mgr.release(b"nonexistent")); + } + + #[test] + fn test_barrier_manager_lease_association() { + let mut mgr = BarrierManager::new(); + let name = b"test-barrier".to_vec(); + let key = barrier_key(&name); + + mgr.create(name.clone(), key.clone(), 42); + assert_eq!(mgr.get_lease_id(&key), Some(42)); + } + + #[test] + fn test_barrier_manager_release_by_lease() { + let mut mgr = BarrierManager::new(); + let name1 = b"barrier-1".to_vec(); + let name2 = b"barrier-2".to_vec(); + let key1 = barrier_key(&name1); + let key2 = barrier_key(&name2); + + mgr.create(name1.clone(), key1.clone(), 42); + mgr.create(name2.clone(), key2.clone(), 42); + + assert_eq!(mgr.barrier_count(), 2); + + let released = mgr.release_by_lease(42); + assert_eq!(released.len(), 2); + assert_eq!(mgr.barrier_count(), 0); + } + + #[test] + fn test_validate_barrier_name() { + assert!(validate_barrier_name(b"valid-name").is_ok()); + assert!(validate_barrier_name(b"").is_err()); + assert!(validate_barrier_name(b"null\x00byte").is_err()); + assert!(validate_barrier_name(b"_aether_reserved").is_err()); + } +} diff --git a/src/lib.rs b/src/lib.rs index 4f7cafe..b8cdf89 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ pub mod api; pub mod auth; +pub mod barrier; pub mod cluster; pub mod config; pub mod election; @@ -7,6 +8,7 @@ pub mod error; pub mod lease; pub mod lock; pub mod proto; +pub mod queue; pub mod raft; pub mod shard; pub mod storage; diff --git a/src/main.rs b/src/main.rs index cdb252f..4c85735 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,25 +10,29 @@ use tracing_subscriber::{EnvFilter, Layer, fmt, layer::SubscriberExt, util::Subs use aether::api::health::HealthStatus; use aether::api::metrics::{MetricsLayer, MetricsRegistry}; use aether::api::{ - AuthService, ClusterService, ElectionService, KvService, LeaseService, LockService, - MaintenanceService, ShardService, WatchService, + AuthService, BarrierService, ClusterService, ElectionService, KvService, LeaseService, + LockService, MaintenanceService, QueueService, ShardService, WatchService, }; use aether::auth::AuthLayer; +use aether::barrier::BarrierManager; use aether::cluster::AlarmManager; use aether::config::{AetherConfig, LogConfig}; use aether::election::ElectionManager; use aether::lease::{LeaseManager, LeaseStore}; use aether::lock::LockManager; use aether::proto::aether_auth_server::AetherAuthServer; +use aether::proto::aether_barrier_server::AetherBarrierServer; use aether::proto::aether_cluster_server::AetherClusterServer; use aether::proto::aether_election_server::AetherElectionServer; use aether::proto::aether_kv_server::AetherKvServer; use aether::proto::aether_lease_server::AetherLeaseServer; use aether::proto::aether_lock_server::AetherLockServer; use aether::proto::aether_maintenance_server::AetherMaintenanceServer; +use aether::proto::aether_queue_server::AetherQueueServer; use aether::proto::aether_shard_server::AetherShardServer; use aether::proto::aether_watch_server::AetherWatchServer; use aether::proto::raft_rpc::raft_rpc_server::RaftRpcServer; +use aether::queue::QueueManager; use aether::raft::node; use aether::raft::raftrs_handle::RaftRsHandle; use aether::raft::raftrs_store::RaftRsStore; @@ -311,6 +315,13 @@ async fn main() -> anyhow::Result<()> { let election_manager = Arc::new(Mutex::new(ElectionManager::new())); let election_manager_for_api = election_manager.clone(); + // Create barrier manager (shared between state machine and API layer) + let barrier_manager = Arc::new(Mutex::new(BarrierManager::new())); + let barrier_manager_for_api = barrier_manager.clone(); + + // Create queue manager (shared between state machine and API layer) + let queue_manager = Arc::new(Mutex::new(QueueManager::new())); + let state_machine = Arc::new(Mutex::new(AetherStateMachine::new( watch_tx.clone(), storage.clone(), @@ -321,6 +332,8 @@ async fn main() -> anyhow::Result<()> { shard_manager, lock_manager, election_manager, + barrier_manager, + queue_manager, ))); // auth_enabled is cloned above; keep a reference for ClusterService below @@ -381,6 +394,7 @@ async fn main() -> anyhow::Result<()> { let watch_manager = WatchManager::new(watch_tx); let watch_manager_for_metrics = watch_manager.clone(); let watch_manager_for_election = watch_manager.clone(); + let watch_manager_for_barrier = watch_manager.clone(); let watch_service = WatchService::new( watch_manager, auth_enabled.clone(), @@ -399,6 +413,7 @@ async fn main() -> anyhow::Result<()> { let auth_enabled_for_api = auth_enabled.clone(); let storage_for_maintenance = storage.clone(); let storage_for_election = storage.clone(); + let storage_for_queue = storage.clone(); let kv_service = KvService::new( storage, raft_handle.clone(), @@ -445,6 +460,15 @@ async fn main() -> anyhow::Result<()> { watch_manager_for_election, ); + let barrier_service = BarrierService::new( + raft_handle.clone(), + config.node_id, + barrier_manager_for_api, + watch_manager_for_barrier, + ); + + let queue_service = QueueService::new(raft_handle.clone(), config.node_id, storage_for_queue); + let alarm_manager = Arc::new(AlarmManager::new()); let maintenance_service = MaintenanceService::new( raft_handle.clone(), @@ -571,6 +595,8 @@ async fn main() -> anyhow::Result<()> { .add_service(AetherLeaseServer::new(lease_service)) .add_service(AetherLockServer::new(lock_service)) .add_service(AetherElectionServer::new(election_service)) + .add_service(AetherBarrierServer::new(barrier_service)) + .add_service(AetherQueueServer::new(queue_service)) .add_service(AetherClusterServer::new(cluster_service)) .add_service(AetherMaintenanceServer::new(maintenance_service)) .add_service(AetherShardServer::new(shard_service)) diff --git a/src/queue/mod.rs b/src/queue/mod.rs new file mode 100644 index 0000000..eff2454 --- /dev/null +++ b/src/queue/mod.rs @@ -0,0 +1,241 @@ +use std::collections::HashMap; + +use crate::storage::{RocksStorage, StorageEngine}; + +/// Key prefix for queue item keys in storage. +/// Items are stored as `_aether_queue//`. +pub const QUEUE_KEY_PREFIX: &[u8] = b"_aether_queue/"; + +/// Key prefix for queue metadata (next sequence number). +pub const QUEUE_META_PREFIX: &[u8] = b"_aether_queue_meta/"; + +/// Reserved prefix for system use. +pub const RESERVED_PREFIX: &[u8] = b"_aether_"; + +/// Maximum allowed queue name length (1 KB). +pub const MAX_QUEUE_NAME_LEN: usize = 1024; + +/// In-memory queue manager. Owned by the state machine. +/// +/// Tracks queue metadata (next sequence number per queue). +/// Actual queue items are stored in the KV store with sequential keys. +pub struct QueueManager { + /// Next sequence number per queue name. + next_seq: HashMap, u64>, +} + +impl Default for QueueManager { + fn default() -> Self { + Self::new() + } +} + +impl QueueManager { + pub fn new() -> Self { + Self { + next_seq: HashMap::new(), + } + } + + /// Restore queue state from persistent storage. + /// Scans the queue meta prefix to rebuild sequence counters. + pub fn restore(&mut self, storage: &RocksStorage) -> Result<(), crate::error::StorageError> { + let entries = storage.scan(QUEUE_META_PREFIX, usize::MAX)?; + + self.next_seq.clear(); + + for entry in entries { + let meta_key = &entry.key; + let value = entry.value; + + // Extract queue name from meta key + let name = match queue_meta_name(meta_key) { + Some(n) => n, + None => { + tracing::warn!( + key = %String::from_utf8_lossy(meta_key), + "invalid queue meta key, skipping" + ); + continue; + } + }; + + if value.len() < 8 { + tracing::warn!( + key = %String::from_utf8_lossy(meta_key), + "invalid queue meta value, skipping" + ); + continue; + } + + let seq_bytes: [u8; 8] = match value[..8].try_into() { + Ok(bytes) => bytes, + Err(_) => continue, + }; + let seq = u64::from_be_bytes(seq_bytes); + + self.next_seq.insert(name.into(), seq); + } + + Ok(()) + } + + /// Get and increment the next sequence number for a queue. + /// Returns the current sequence number. + pub fn next_seq(&mut self, name: &[u8]) -> u64 { + let seq = self.next_seq.get(name).copied().unwrap_or(1); + self.next_seq.insert(name.into(), seq + 1); + seq + } + + /// Get the current sequence number without incrementing. + pub fn peek_seq(&self, name: &[u8]) -> u64 { + self.next_seq.get(name).copied().unwrap_or(1) + } + + /// Get the number of tracked queues. + pub fn queue_count(&self) -> usize { + self.next_seq.len() + } +} + +/// Generate a queue item key: `_aether_queue//`. +pub fn queue_item_key(name: &[u8], seq: u64) -> Vec { + let seq_bytes = seq.to_be_bytes(); + let mut key = Vec::with_capacity(QUEUE_KEY_PREFIX.len() + name.len() + 1 + 8); + key.extend_from_slice(QUEUE_KEY_PREFIX); + key.extend_from_slice(name); + key.push(b'/'); + key.extend_from_slice(&seq_bytes); + key +} + +/// Generate the prefix for scanning all items in a queue: `_aether_queue//`. +pub fn queue_scan_prefix(name: &[u8]) -> Vec { + let mut prefix = Vec::with_capacity(QUEUE_KEY_PREFIX.len() + name.len() + 1); + prefix.extend_from_slice(QUEUE_KEY_PREFIX); + prefix.extend_from_slice(name); + prefix.push(b'/'); + prefix +} + +/// Generate the meta key for a queue: `_aether_queue_meta/`. +pub fn queue_meta_key(name: &[u8]) -> Vec { + let mut key = Vec::with_capacity(QUEUE_META_PREFIX.len() + name.len()); + key.extend_from_slice(QUEUE_META_PREFIX); + key.extend_from_slice(name); + key +} + +/// Extract the queue name from a meta key. +fn queue_meta_name(key: &[u8]) -> Option<&[u8]> { + if key.starts_with(QUEUE_META_PREFIX) { + Some(&key[QUEUE_META_PREFIX.len()..]) + } else { + None + } +} + +/// Validate queue name. +pub fn validate_queue_name(name: &[u8]) -> Result<(), &'static str> { + if name.is_empty() { + return Err("queue name must not be empty"); + } + if name.len() > MAX_QUEUE_NAME_LEN { + return Err("queue name too long"); + } + if name.contains(&0) { + return Err("queue name must not contain null bytes"); + } + if name.contains(&b'/') { + return Err("queue name must not contain '/'"); + } + if name.starts_with(RESERVED_PREFIX) { + return Err("queue name must not start with reserved prefix '_aether_'"); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_queue_item_key_generation() { + let name = b"my-queue"; + let key = queue_item_key(name, 1); + let expected_seq = 1u64.to_be_bytes(); + assert!(key.starts_with(b"_aether_queue/my-queue/")); + assert_eq!(&key[key.len() - 8..], expected_seq); + } + + #[test] + fn test_queue_scan_prefix() { + let prefix = queue_scan_prefix(b"my-queue"); + assert_eq!(prefix, b"_aether_queue/my-queue/"); + } + + #[test] + fn test_queue_meta_key() { + let key = queue_meta_key(b"my-queue"); + assert_eq!(key, b"_aether_queue_meta/my-queue"); + } + + #[test] + fn test_queue_meta_name_extraction() { + let key = b"_aether_queue_meta/my-queue"; + let name = queue_meta_name(key); + assert_eq!(name, Some(b"my-queue".as_slice())); + } + + #[test] + fn test_queue_manager_next_seq() { + let mut mgr = QueueManager::new(); + let name = b"test-queue"; + + assert_eq!(mgr.next_seq(name), 1); + assert_eq!(mgr.next_seq(name), 2); + assert_eq!(mgr.next_seq(name), 3); + } + + #[test] + fn test_queue_manager_peek_seq() { + let mut mgr = QueueManager::new(); + let name = b"test-queue"; + + assert_eq!(mgr.peek_seq(name), 1); + mgr.next_seq(name); + assert_eq!(mgr.peek_seq(name), 2); + } + + #[test] + fn test_queue_manager_multiple_queues() { + let mut mgr = QueueManager::new(); + + assert_eq!(mgr.next_seq(b"queue-a"), 1); + assert_eq!(mgr.next_seq(b"queue-b"), 1); + assert_eq!(mgr.next_seq(b"queue-a"), 2); + assert_eq!(mgr.next_seq(b"queue-b"), 2); + } + + #[test] + fn test_validate_queue_name() { + assert!(validate_queue_name(b"valid-name").is_ok()); + assert!(validate_queue_name(b"").is_err()); + assert!(validate_queue_name(b"null\x00byte").is_err()); + assert!(validate_queue_name(b"_aether_reserved").is_err()); + assert!(validate_queue_name(b"has/slash").is_err()); + } + + #[test] + fn test_queue_item_key_ordering() { + let name = b"test"; + let key1 = queue_item_key(name, 1); + let key2 = queue_item_key(name, 2); + let key10 = queue_item_key(name, 10); + + // Big-endian encoding ensures lexicographic order matches logical order + assert!(key1 < key2); + assert!(key2 < key10); + } +} diff --git a/src/raft/mod.rs b/src/raft/mod.rs index 006413a..1213b62 100644 --- a/src/raft/mod.rs +++ b/src/raft/mod.rs @@ -95,6 +95,14 @@ pub enum RaftRequest { }, /// Resign from leadership in an election ElectionResign { leader_key: Vec }, + /// Create and hold a barrier + BarrierCreate { name: Vec, lease_id: i64 }, + /// Release a barrier + BarrierRelease { name: Vec }, + /// Enqueue an item to a named queue + QueueEnqueue { name: Vec, value: Vec }, + /// Dequeue the front item from a named queue + QueueDequeue { name: Vec }, } /// Raft response types @@ -153,6 +161,23 @@ pub enum RaftResponse { ElectionResign {}, /// Election resign requested but leader key not found ElectionResignNotFound {}, + BarrierCreate { + key: Vec, + }, + /// Barrier already held by another caller + BarrierAlreadyHeld { + current_key: Vec, + }, + BarrierRelease {}, + QueueEnqueue { + key: Vec, + }, + QueueDequeue { + key: Vec, + value: Vec, + }, + /// Queue is empty + QueueDequeueEmpty {}, /// Storage error during apply Error { message: String, diff --git a/src/raft/state_machine.rs b/src/raft/state_machine.rs index 91cc486..5b857e4 100644 --- a/src/raft/state_machine.rs +++ b/src/raft/state_machine.rs @@ -4,11 +4,13 @@ use std::sync::{Arc, Mutex}; use rocksdb::WriteBatch; -use super::{KeyValue, RaftRequest, RaftResponse, WatchEvent}; +use super::{KeyValue, RaftRequest, RaftResponse, WatchEvent, WatchEventType}; use crate::auth::AuthCache; +use crate::barrier::BarrierManager; use crate::election::ElectionManager; use crate::lease::{LeaseManager, LeaseStore}; use crate::lock::LockManager; +use crate::queue::QueueManager; use crate::shard::manager::ShardManager; use crate::storage::mvcc::{ KeyIndex, MvccValue, encode_mvcc_key, load_key_indexes, save_global_revision, @@ -84,6 +86,10 @@ pub struct AetherStateMachine { pub lock_manager: Arc>, /// In-memory election manager (shared with API layer). pub election_manager: Arc>, + /// In-memory barrier manager (shared with API layer). + pub barrier_manager: Arc>, + /// In-memory queue manager (shared with API layer). + pub queue_manager: Arc>, } /// Captures pre-mutation in-memory state for keys touched during an apply. @@ -204,6 +210,8 @@ impl AetherStateMachine { shard_manager: Arc>, lock_manager: Arc>, election_manager: Arc>, + barrier_manager: Arc>, + queue_manager: Arc>, ) -> Self { let key_indexes = load_key_indexes(storage.db(), storage.mvcc_cf()) .expect("failed to load key indexes from mvcc CF"); @@ -226,6 +234,22 @@ impl AetherStateMachine { } } + // Restore barrier state from persistent storage + { + let mut mgr = barrier_manager.lock().unwrap(); + if let Err(e) = mgr.restore(&storage) { + tracing::error!(error = %e, "failed to restore barrier state from storage"); + } + } + + // Restore queue state from persistent storage + { + let mut mgr = queue_manager.lock().unwrap(); + if let Err(e) = mgr.restore(&storage) { + tracing::error!(error = %e, "failed to restore queue state from storage"); + } + } + Self { last_applied: 0, storage, @@ -240,6 +264,8 @@ impl AetherStateMachine { shard_manager, lock_manager, election_manager, + barrier_manager, + queue_manager, } } @@ -1031,6 +1057,14 @@ impl AetherStateMachine { RaftRequest::ElectionResign { leader_key } => { self.apply_election_resign(leader_key, revision) } + RaftRequest::BarrierCreate { name, lease_id } => { + self.apply_barrier_create(name, lease_id, revision) + } + RaftRequest::BarrierRelease { name } => self.apply_barrier_release(name, revision), + RaftRequest::QueueEnqueue { name, value } => { + self.apply_queue_enqueue(name, value, revision) + } + RaftRequest::QueueDequeue { name } => self.apply_queue_dequeue(name, revision), } } @@ -1151,6 +1185,9 @@ impl AetherStateMachine { // Release all elections associated with this lease self.apply_election_release_by_lease(id); + // Release all barriers associated with this lease + self.apply_barrier_release_by_lease(id); + // If no user-data keys and lease already gone, this is an idempotent no-op. if key_values.is_empty() { let mgr = self.lease_manager.lock().unwrap(); @@ -2013,6 +2050,262 @@ impl AetherStateMachine { ); } + fn apply_barrier_create( + &mut self, + name: Vec, + lease_id: i64, + _revision: u64, + ) -> RaftResponse { + // Check if barrier already exists + { + let mgr = self.barrier_manager.lock().unwrap(); + if mgr.is_held(&name) { + let existing_key = mgr.get_key(&name).unwrap_or_default().to_vec(); + return RaftResponse::BarrierAlreadyHeld { + current_key: existing_key, + }; + } + } + + let key = crate::barrier::barrier_key(&name); + let value = crate::barrier::BarrierManager::encode_value(&name, lease_id); + + // Write to KV store + let mut batch = WriteBatch::default(); + batch.put_cf(self.storage.default_cf(), &key, &value); + if let Err(e) = self.storage.db().write(batch) { + return RaftResponse::Error { + message: format!("failed to write barrier: {e}"), + }; + } + + // Update in-memory state + { + let mut mgr = self.barrier_manager.lock().unwrap(); + mgr.create(name, key.clone(), lease_id); + } + + tracing::debug!( + barrier = %String::from_utf8_lossy(&key), + lease_id = lease_id, + "barrier created" + ); + + RaftResponse::BarrierCreate { key } + } + + fn apply_barrier_release(&mut self, name: Vec, _revision: u64) -> RaftResponse { + let key = crate::barrier::barrier_key(&name); + + // Check if barrier exists + { + let mgr = self.barrier_manager.lock().unwrap(); + if !mgr.is_held(&name) { + return RaftResponse::BarrierRelease {}; + } + } + + // Delete from KV store + let mut batch = WriteBatch::default(); + batch.delete_cf(self.storage.default_cf(), &key); + if let Err(e) = self.storage.db().write(batch) { + return RaftResponse::Error { + message: format!("failed to delete barrier: {e}"), + }; + } + + // Update in-memory state + { + let mut mgr = self.barrier_manager.lock().unwrap(); + mgr.release(&key); + } + + // Emit watch event so waiting clients can be notified + let event = WatchEvent { + event_type: WatchEventType::Delete, + kv: KeyValue { + key, + value: Vec::new(), + create_revision: 0, + mod_revision: 0, + version: 0, + lease: 0, + }, + prev_kv: None, + }; + let _ = self.watch_tx.send(event); + + tracing::debug!( + barrier = %String::from_utf8_lossy(&name), + "barrier released" + ); + + RaftResponse::BarrierRelease {} + } + + fn apply_barrier_release_by_lease(&mut self, lease_id: i64) { + let keys_to_release: Vec> = { + let mgr = self.barrier_manager.lock().unwrap(); + mgr.get_keys_by_lease(lease_id) + }; + + if keys_to_release.is_empty() { + return; + } + + let mut batch = WriteBatch::default(); + for key in &keys_to_release { + batch.delete_cf(self.storage.default_cf(), key); + } + + if let Err(e) = self.storage.db().write(batch) { + tracing::error!(error = %e, lease_id = lease_id, "failed to delete barrier keys for lease"); + return; + } + + { + let mut mgr = self.barrier_manager.lock().unwrap(); + mgr.release_by_lease(lease_id); + } + + // Emit watch events for each released barrier + let count = keys_to_release.len(); + for key in keys_to_release { + let event = WatchEvent { + event_type: WatchEventType::Delete, + kv: KeyValue { + key, + value: Vec::new(), + create_revision: 0, + mod_revision: 0, + version: 0, + lease: 0, + }, + prev_kv: None, + }; + let _ = self.watch_tx.send(event); + } + + tracing::debug!( + lease_id = lease_id, + count = count, + "released barriers for lease" + ); + } + + fn apply_queue_enqueue( + &mut self, + name: Vec, + value: Vec, + _revision: u64, + ) -> RaftResponse { + // Get next sequence number and meta key + let (seq, meta_key) = { + let mut mgr = self.queue_manager.lock().unwrap(); + let seq = mgr.next_seq(&name); + let meta_key = crate::queue::queue_meta_key(&name); + (seq, meta_key) + }; + + let item_key = crate::queue::queue_item_key(&name, seq); + // Persist next sequence number (seq + 1), not current seq. + // On restore, this value is loaded as the next seq to use. + let next_seq = seq.checked_add(1).unwrap_or(0); + let next_seq_bytes = next_seq.to_be_bytes(); + + // Write item and update meta atomically + let mut batch = WriteBatch::default(); + batch.put_cf(self.storage.default_cf(), &item_key, &value); + batch.put_cf(self.storage.default_cf(), &meta_key, next_seq_bytes); + if let Err(e) = self.storage.db().write(batch) { + return RaftResponse::Error { + message: format!("failed to enqueue: {e}"), + }; + } + + // Emit watch event so waiting clients can be notified. + // Use empty value in watch event to avoid cloning potentially large values. + // Clients that need the value can fetch it from storage using the key. + let event = WatchEvent { + event_type: WatchEventType::Put, + kv: KeyValue { + key: item_key.clone(), + value: Vec::new(), + create_revision: 0, + mod_revision: 0, + version: 0, + lease: 0, + }, + prev_kv: None, + }; + let _ = self.watch_tx.send(event); + + tracing::debug!( + queue = %String::from_utf8_lossy(&name), + seq = seq, + "item enqueued" + ); + + RaftResponse::QueueEnqueue { key: item_key } + } + + fn apply_queue_dequeue(&mut self, name: Vec, _revision: u64) -> RaftResponse { + // Find the front item by scanning with the queue prefix + let prefix = crate::queue::queue_scan_prefix(&name); + let entries = match self.storage.scan(&prefix, 1) { + Ok(entries) => entries, + Err(e) => { + return RaftResponse::Error { + message: format!("failed to scan queue: {e}"), + }; + } + }; + + if entries.is_empty() { + return RaftResponse::QueueDequeueEmpty {}; + } + + let front = &entries[0]; + let item_key = front.key.clone(); + let item_value = front.value.clone(); + + // Delete the item from KV store + let mut batch = WriteBatch::default(); + batch.delete_cf(self.storage.default_cf(), &item_key); + if let Err(e) = self.storage.db().write(batch) { + return RaftResponse::Error { + message: format!("failed to dequeue: {e}"), + }; + } + + // Emit watch event for the deleted item. + // Use empty value in watch event to avoid cloning potentially large values. + let event = WatchEvent { + event_type: WatchEventType::Delete, + kv: KeyValue { + key: item_key.clone(), + value: Vec::new(), + create_revision: 0, + mod_revision: 0, + version: 0, + lease: 0, + }, + prev_kv: None, + }; + let _ = self.watch_tx.send(event); + + tracing::debug!( + queue = %String::from_utf8_lossy(&name), + key = %String::from_utf8_lossy(&item_key), + "item dequeued" + ); + + RaftResponse::QueueDequeue { + key: item_key, + value: item_value, + } + } + /// Maximum allowed snapshot size (512 MiB). Prevents OOM when the database /// is larger than available memory. The serialized snapshot is typically /// 1-2x the raw CF data size, so this limits total memory to ~1.5 GiB. @@ -2171,12 +2464,30 @@ impl AetherStateMachine { } } + // Rebuild barrier manager from persisted barrier data. + { + let mut mgr = self.barrier_manager.lock().unwrap(); + if let Err(e) = mgr.restore(&self.storage) { + tracing::warn!(error = %e, "failed to restore barrier manager from snapshot"); + } + } + + // Rebuild queue manager from persisted queue data. + { + let mut mgr = self.queue_manager.lock().unwrap(); + if let Err(e) = mgr.restore(&self.storage) { + tracing::warn!(error = %e, "failed to restore queue manager from snapshot"); + } + } + self.last_applied = applied_index; tracing::info!( applied_index, keys = self.key_metas.len(), locks = self.lock_manager.lock().unwrap().lock_count(), elections = self.election_manager.lock().unwrap().election_count(), + barriers = self.barrier_manager.lock().unwrap().barrier_count(), + queues = self.queue_manager.lock().unwrap().queue_count(), "restored from snapshot" ); Ok(()) @@ -2262,8 +2573,10 @@ impl AetherStateMachine { #[cfg(test)] mod tests { use super::*; + use crate::barrier::BarrierManager; use crate::election::ElectionManager; use crate::lease::now_millis; + use crate::queue::QueueManager; use crate::raft::{self, WatchEventType}; use tempfile::tempdir; @@ -2279,6 +2592,8 @@ mod tests { let shard_manager = Arc::new(Mutex::new(ShardManager::new())); let lock_manager = Arc::new(Mutex::new(LockManager::new())); let election_manager = Arc::new(Mutex::new(ElectionManager::new())); + let barrier_manager = Arc::new(Mutex::new(BarrierManager::new())); + let queue_manager = Arc::new(Mutex::new(QueueManager::new())); let sm = AetherStateMachine::new( tx.clone(), storage.clone(), @@ -2289,6 +2604,8 @@ mod tests { shard_manager, lock_manager, election_manager, + barrier_manager, + queue_manager, ); (dir, storage, sm) } @@ -2459,6 +2776,8 @@ mod tests { let shard_manager = Arc::new(Mutex::new(ShardManager::new())); let lock_manager = Arc::new(Mutex::new(LockManager::new())); let election_manager = Arc::new(Mutex::new(ElectionManager::new())); + let barrier_manager = Arc::new(Mutex::new(BarrierManager::new())); + let queue_manager = Arc::new(Mutex::new(QueueManager::new())); let mut sm = AetherStateMachine::new( tx, storage, @@ -2469,6 +2788,8 @@ mod tests { shard_manager, lock_manager, election_manager, + barrier_manager, + queue_manager, ); sm.apply_request( @@ -2500,6 +2821,8 @@ mod tests { let shard_manager = Arc::new(Mutex::new(ShardManager::new())); let lock_manager = Arc::new(Mutex::new(LockManager::new())); let election_manager = Arc::new(Mutex::new(ElectionManager::new())); + let barrier_manager = Arc::new(Mutex::new(BarrierManager::new())); + let queue_manager = Arc::new(Mutex::new(QueueManager::new())); let mut sm = AetherStateMachine::new( tx, storage, @@ -2510,6 +2833,8 @@ mod tests { shard_manager, lock_manager, election_manager, + barrier_manager, + queue_manager, ); sm.apply_request( @@ -2762,6 +3087,8 @@ mod tests { let shard_manager = Arc::new(Mutex::new(ShardManager::new())); let lock_manager = Arc::new(Mutex::new(LockManager::new())); let election_manager = Arc::new(Mutex::new(ElectionManager::new())); + let barrier_manager = Arc::new(Mutex::new(BarrierManager::new())); + let queue_manager = Arc::new(Mutex::new(QueueManager::new())); let mut sm = AetherStateMachine::new( tx, storage, @@ -2772,6 +3099,8 @@ mod tests { shard_manager, lock_manager, election_manager, + barrier_manager, + queue_manager, ); sm.apply_request( @@ -3628,6 +3957,8 @@ mod tests { let shard_manager2 = Arc::new(Mutex::new(ShardManager::new())); let lock_manager2 = Arc::new(Mutex::new(LockManager::new())); let election_manager2 = Arc::new(Mutex::new(ElectionManager::new())); + let barrier_manager2 = Arc::new(Mutex::new(BarrierManager::new())); + let queue_manager2 = Arc::new(Mutex::new(QueueManager::new())); let mut sm2 = AetherStateMachine::new( tx2, storage.clone(), @@ -3638,6 +3969,8 @@ mod tests { shard_manager2, lock_manager2, election_manager2, + barrier_manager2, + queue_manager2, ); // Write different data to sm2 to verify it gets overwritten. @@ -3703,6 +4036,8 @@ mod tests { let shard_manager2 = Arc::new(Mutex::new(ShardManager::new())); let lock_manager2 = Arc::new(Mutex::new(LockManager::new())); let election_manager2 = Arc::new(Mutex::new(ElectionManager::new())); + let barrier_manager2 = Arc::new(Mutex::new(BarrierManager::new())); + let queue_manager2 = Arc::new(Mutex::new(QueueManager::new())); let mut sm2 = AetherStateMachine::new( tx2, storage.clone(), @@ -3713,6 +4048,8 @@ mod tests { shard_manager2, lock_manager2, election_manager2, + barrier_manager2, + queue_manager2, ); sm2.restore_snapshot(&snapshot_data, 2).unwrap(); @@ -3768,6 +4105,8 @@ mod tests { let shard_manager2 = Arc::new(Mutex::new(ShardManager::new())); let lock_manager2 = Arc::new(Mutex::new(LockManager::new())); let election_manager2 = Arc::new(Mutex::new(ElectionManager::new())); + let barrier_manager2 = Arc::new(Mutex::new(BarrierManager::new())); + let queue_manager2 = Arc::new(Mutex::new(QueueManager::new())); let mut sm2 = AetherStateMachine::new( tx2, storage.clone(), @@ -3778,6 +4117,8 @@ mod tests { shard_manager2, lock_manager2, election_manager2, + barrier_manager2, + queue_manager2, ); // Verify lock exists in fresh state machine (from startup restore). @@ -3853,6 +4194,8 @@ mod tests { let shard_manager2 = Arc::new(Mutex::new(ShardManager::new())); let lock_manager2 = Arc::new(Mutex::new(LockManager::new())); let election_manager2 = Arc::new(Mutex::new(ElectionManager::new())); + let barrier_manager2 = Arc::new(Mutex::new(BarrierManager::new())); + let queue_manager2 = Arc::new(Mutex::new(QueueManager::new())); let mut sm2 = AetherStateMachine::new( tx2, storage.clone(), @@ -3863,6 +4206,8 @@ mod tests { shard_manager2, lock_manager2, election_manager2, + barrier_manager2, + queue_manager2, ); // Restore from snapshot.