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.