use std::{future::Future, pin::Pin, sync::Arc};
pub trait ClusterRoleGate: Send + Sync {
fn prevent_role_change(&self) -> bool;
fn allow_role_change(&self);
}
pub struct AlwaysAllowGate;
impl ClusterRoleGate for AlwaysAllowGate {
fn prevent_role_change(&self) -> bool {
true
}
fn allow_role_change(&self) {}
}
pub struct RoleChangeGuard<'a> {
gate: Option<&'a dyn ClusterRoleGate>,
}
impl Drop for RoleChangeGuard<'_> {
fn drop(&mut self) {
if let Some(gate) = self.gate {
gate.allow_role_change();
}
}
}
pub trait StoreCommitFace: Send + Sync {
fn wait_for_commit(&self) -> bool;
fn commit_aof(&self, db_id: i64) -> wkv::Result<()>;
fn flush_database<'a>(
&'a self,
unsafe_truncate_log: bool,
db_id: i64,
) -> Pin<Box<dyn Future<Output = wkv::Result<()>> + Send + 'a>>;
}
pub struct StoreApi {
store: Arc<dyn StoreCommitFace>,
cluster_gate: Option<Arc<dyn ClusterRoleGate>>,
is_replica: bool,
}
impl StoreApi {
pub fn new(
store: Arc<dyn StoreCommitFace>,
cluster_gate: Option<Arc<dyn ClusterRoleGate>>,
) -> Self {
Self {
store,
cluster_gate,
is_replica: false,
}
}
pub fn set_replica(&mut self, is_replica: bool) {
self.is_replica = is_replica;
}
fn prevent_role_change(&self) -> Option<RoleChangeGuard<'_>> {
let Some(gate) = self.cluster_gate.as_deref() else {
return Some(RoleChangeGuard { gate: None });
};
if !gate.prevent_role_change() {
return None;
}
Some(RoleChangeGuard { gate: Some(gate) })
}
pub fn wait_for_commit(&self) -> bool {
let _guard = match self.prevent_role_change() {
Some(guard) => guard,
None => return false,
};
if self.is_replica {
return false;
}
self.store.wait_for_commit()
}
pub fn commit_aof(&self, db_id: i64) -> wkv::Result<bool> {
let _guard = match self.prevent_role_change() {
Some(guard) => guard,
None => return Ok(false),
};
if self.is_replica {
return Ok(false);
}
self.store.commit_aof(db_id).map(|()| true)
}
pub async fn flush_db(&self, db_id: i64, unsafe_truncate_log: bool) -> wkv::Result<bool> {
let _guard = match self.prevent_role_change() {
Some(guard) => guard,
None => return Ok(false),
};
if self.is_replica {
return Ok(false);
}
self
.store
.flush_database(unsafe_truncate_log, db_id)
.await
.map(|()| true)
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicBool, Ordering};
use compio::runtime::Runtime;
use super::*;
struct RecordingGate {
allowed: AtomicBool,
prevented: AtomicBool,
released: AtomicBool,
}
impl RecordingGate {
fn new() -> Arc<Self> {
Arc::new(Self {
allowed: AtomicBool::new(true),
prevented: AtomicBool::new(false),
released: AtomicBool::new(false),
})
}
}
impl ClusterRoleGate for RecordingGate {
fn prevent_role_change(&self) -> bool {
self.prevented.store(true, Ordering::SeqCst);
self.allowed.load(Ordering::SeqCst)
}
fn allow_role_change(&self) {
self.released.store(true, Ordering::SeqCst);
}
}
struct MockStore {
waited: AtomicBool,
committed: AtomicBool,
flushed: AtomicBool,
}
impl MockStore {
fn new() -> Arc<Self> {
Arc::new(Self {
waited: AtomicBool::new(false),
committed: AtomicBool::new(false),
flushed: AtomicBool::new(false),
})
}
}
impl StoreCommitFace for MockStore {
fn wait_for_commit(&self) -> bool {
self.waited.store(true, Ordering::SeqCst);
true
}
fn commit_aof(&self, _db_id: i64) -> wkv::Result<()> {
self.committed.store(true, Ordering::SeqCst);
Ok(())
}
fn flush_database<'a>(
&'a self,
_unsafe_truncate_log: bool,
_db_id: i64,
) -> Pin<Box<dyn Future<Output = wkv::Result<()>> + Send + 'a>> {
Box::pin(async {
self.flushed.store(true, Ordering::SeqCst);
Ok(())
})
}
}
#[test]
fn commit_gates_allow_and_reject() {
let store = MockStore::new();
let gate = RecordingGate::new();
let api = StoreApi::new(store.clone(), Some(gate.clone()));
assert!(api.commit_aof(0).expect("提交成功"));
assert!(gate.prevented.load(Ordering::SeqCst));
assert!(gate.released.load(Ordering::SeqCst));
assert!(store.committed.load(Ordering::SeqCst));
gate.allowed.store(false, Ordering::SeqCst);
store.committed.store(false, Ordering::SeqCst);
assert!(!api.commit_aof(0).expect("路径闭环"));
assert!(!store.committed.load(Ordering::SeqCst));
}
#[test]
fn replica_state_ignores_commits() {
let store = MockStore::new();
let gate = RecordingGate::new();
let mut api = StoreApi::new(store.clone(), Some(gate));
api.set_replica(true);
assert!(!api.commit_aof(0).expect("路径闭环"));
assert!(!api.wait_for_commit());
assert!(!store.committed.load(Ordering::SeqCst));
assert!(!store.waited.load(Ordering::SeqCst));
}
#[test]
fn non_cluster_always_allowed() {
let store = MockStore::new();
let api = StoreApi::new(store.clone(), None);
assert!(api.wait_for_commit());
assert!(api.commit_aof(0).expect("提交成功"));
assert!(store.waited.load(Ordering::SeqCst));
assert!(store.committed.load(Ordering::SeqCst));
}
#[test]
fn flush_db_respects_gate() {
let store = MockStore::new();
let gate = RecordingGate::new();
let api = StoreApi::new(store.clone(), Some(gate.clone()));
let rt = Runtime::new().expect("运行时");
rt.block_on(async {
assert!(api.flush_db(0, false).await.expect("清库成功"));
assert!(store.flushed.load(Ordering::SeqCst));
gate.allowed.store(false, Ordering::SeqCst);
store.flushed.store(false, Ordering::SeqCst);
assert!(!api.flush_db(0, false).await.expect("路径闭环"));
assert!(!store.flushed.load(Ordering::SeqCst));
});
}
}