diff --git a/src/meta-srv/src/handler/collect_stats_handler.rs b/src/meta-srv/src/handler/collect_stats_handler.rs index d444498798..e3ae7ebe91 100644 --- a/src/meta-srv/src/handler/collect_stats_handler.rs +++ b/src/meta-srv/src/handler/collect_stats_handler.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::cmp::Ordering; +use std::sync::Arc; use api::v1::meta::{HeartbeatRequest, Role}; use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; @@ -24,6 +25,7 @@ use common_meta::rpc::store::PutRequest; use common_telemetry::{error, info, warn}; use dashmap::DashMap; use snafu::ResultExt; +use tokio::sync::Mutex; use crate::error::{self, Result}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; @@ -70,7 +72,7 @@ impl EpochStats { const DEFAULT_FLUSH_STATS_FACTOR: usize = 3; pub struct CollectStatsHandler { - stats_cache: DashMap, + stats_cache: DashMap>>, flush_stats_factor: usize, } @@ -106,10 +108,16 @@ impl HeartbeatHandler for CollectStatsHandler { }; let key = current_stat.stat_key(); - let mut entry = self.stats_cache.entry(key).or_default(); + let state = { + let entry = self + .stats_cache + .entry(key) + .or_insert_with(|| Arc::new(Mutex::new(EpochStats::default()))); + Arc::clone(entry.value()) + }; + let mut epoch_stats = state.lock().await; let key: Vec = key.into(); - let epoch_stats = entry.value_mut(); let refresh = if let Some(epoch) = epoch_stats.epoch() { match current_stat.node_epoch.cmp(&epoch) { @@ -170,7 +178,7 @@ impl HeartbeatHandler for CollectStatsHandler { } } -async fn rewrite_node_address(ctx: &mut Context, stat: &Stat) { +async fn rewrite_node_address(ctx: &Context, stat: &Stat) { let peer = Peer { id: stat.id, addr: stat.addr.clone(), @@ -217,10 +225,188 @@ async fn rewrite_node_address(ctx: &mut Context, stat: &Stat) { #[cfg(test)] mod tests { + use std::any::Any; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering}; + use std::sync::{Arc, Mutex as StdMutex, mpsc}; + use std::thread; + use std::time::Duration; + use common_meta::datanode::DatanodeStatKey; + use common_meta::error::{Error as MetaError, Result as MetaResult}; + use common_meta::kv_backend::{KvBackend, KvBackendRef, ResettableKvBackend, TxnService}; + use common_meta::rpc::store::{ + BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, + BatchPutRequest, BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, PutResponse, + RangeRequest, RangeResponse, + }; + use tokio::sync::Semaphore; + use tokio::time::{sleep, timeout}; use super::*; use crate::handler::test_utils::TestEnv; + use crate::service::store::cached_kv::LeaderCachedKvBackend; + + struct ControlledKvBackend { + recorded_puts: StdMutex>, + put_entered: Semaphore, + put_release: Semaphore, + block_next_put: AtomicBool, + delay_puts: AtomicBool, + active_puts: AtomicUsize, + max_active_puts: AtomicUsize, + } + + impl ControlledKvBackend { + fn new() -> Self { + Self { + recorded_puts: StdMutex::new(Vec::new()), + put_entered: Semaphore::new(0), + put_release: Semaphore::new(0), + block_next_put: AtomicBool::new(false), + delay_puts: AtomicBool::new(false), + active_puts: AtomicUsize::new(0), + max_active_puts: AtomicUsize::new(0), + } + } + + fn block_next_put(&self) { + self.block_next_put.store(true, AtomicOrdering::Relaxed); + } + + async fn wait_for_blocked_put(&self) { + self.put_entered.acquire().await.unwrap().forget(); + } + + fn release_one_put(&self) { + self.put_release.add_permits(1); + } + + fn set_delay_puts(&self, delay: bool) { + self.delay_puts.store(delay, AtomicOrdering::Relaxed); + } + + fn clear_recorded_puts(&self) { + self.recorded_puts.lock().unwrap().clear(); + } + + fn recorded_puts(&self) -> Vec { + self.recorded_puts.lock().unwrap().clone() + } + + fn max_active_puts(&self) -> usize { + self.max_active_puts.load(AtomicOrdering::Relaxed) + } + + fn start_put(&self, req: &PutRequest) -> PutGuard<'_> { + self.recorded_puts.lock().unwrap().push(req.clone()); + let active = self.active_puts.fetch_add(1, AtomicOrdering::Relaxed) + 1; + self.max_active_puts + .fetch_max(active, AtomicOrdering::Relaxed); + PutGuard { backend: self } + } + } + + struct PutGuard<'a> { + backend: &'a ControlledKvBackend, + } + + impl Drop for PutGuard<'_> { + fn drop(&mut self) { + self.backend + .active_puts + .fetch_sub(1, AtomicOrdering::Relaxed); + } + } + + #[async_trait::async_trait] + impl TxnService for ControlledKvBackend { + type Error = MetaError; + } + + #[async_trait::async_trait] + impl KvBackend for ControlledKvBackend { + fn name(&self) -> &str { + "controlled" + } + + fn as_any(&self) -> &dyn Any { + self + } + + async fn range(&self, _req: RangeRequest) -> MetaResult { + unimplemented!() + } + + async fn put(&self, req: PutRequest) -> MetaResult { + let _guard = self.start_put(&req); + if self.block_next_put.swap(false, AtomicOrdering::Relaxed) { + self.put_entered.add_permits(1); + self.put_release.acquire().await.unwrap().forget(); + } + if self.delay_puts.load(AtomicOrdering::Relaxed) { + sleep(Duration::from_millis(10)).await; + } + Ok(PutResponse::default()) + } + + async fn batch_put(&self, _req: BatchPutRequest) -> MetaResult { + unimplemented!() + } + + async fn batch_get(&self, _req: BatchGetRequest) -> MetaResult { + unimplemented!() + } + + async fn delete_range(&self, _req: DeleteRangeRequest) -> MetaResult { + unimplemented!() + } + + async fn batch_delete(&self, _req: BatchDeleteRequest) -> MetaResult { + unimplemented!() + } + } + + impl ResettableKvBackend for ControlledKvBackend { + fn reset(&self) { + self.clear_recorded_puts(); + } + + fn as_kv_backend_ref(self: Arc) -> KvBackendRef { + self + } + } + + fn stat(node_id: u64, epoch: u64, marker: u64, addr: &str) -> Stat { + Stat { + timestamp_millis: marker as i64, + id: node_id, + addr: addr.to_string(), + region_num: marker, + node_epoch: epoch, + ..Default::default() + } + } + + async fn handle_stat( + handler: Arc, + mut ctx: Context, + stat: Stat, + ) -> Result { + let mut acc = HeartbeatAccumulator { + stat: Some(stat), + ..Default::default() + }; + handler + .handle(&HeartbeatRequest::default(), &mut ctx, &mut acc) + .await + } + + fn use_controlled_address_backend(ctx: &mut Context) -> Arc { + let backend = Arc::new(ControlledKvBackend::new()); + ctx.leader_cached_kv_backend = + Arc::new(LeaderCachedKvBackend::with_always_leader(backend.clone())); + backend + } #[tokio::test] async fn test_handle_datanode_stats() { @@ -250,6 +436,105 @@ mod tests { assert_eq!(handler.flush_stats_factor, val.stats.len()); } + #[test] + fn test_same_datanode_wait_keeps_current_thread_runtime_responsive() { + let (backend_tx, backend_rx) = mpsc::sync_channel(1); + let (timer_tx, timer_rx) = mpsc::sync_channel(1); + let (done_tx, done_rx) = mpsc::sync_channel(1); + + let worker = thread::spawn(move || { + tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap() + .block_on(async move { + let env = TestEnv::new(); + let mut ctx = env.ctx(); + let address_backend = use_controlled_address_backend(&mut ctx); + address_backend.block_next_put(); + backend_tx.send(address_backend.clone()).unwrap(); + + let handler = Arc::new(CollectStatsHandler::default()); + let first = tokio::spawn(handle_stat( + handler.clone(), + ctx.clone(), + stat(101, 1, 1, "dn-101-v1"), + )); + address_backend.wait_for_blocked_put().await; + + let second = + tokio::spawn(handle_stat(handler, ctx, stat(101, 1, 2, "dn-101-v1"))); + tokio::spawn(async move { + sleep(Duration::from_millis(10)).await; + timer_tx.send(()).unwrap(); + }); + tokio::task::yield_now().await; + + first.await.unwrap().unwrap(); + second.await.unwrap().unwrap(); + done_tx.send(()).unwrap(); + }); + }); + + let address_backend = backend_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + let timer_result = timer_rx.recv_timeout(Duration::from_secs(1)); + address_backend.release_one_put(); + timer_result.expect("the current-thread runtime must remain responsive"); + done_rx + .recv_timeout(Duration::from_secs(1)) + .expect("both heartbeat handlers must complete after releasing the write"); + worker.join().unwrap(); + } + + #[tokio::test] + async fn test_concurrent_flush_persists_every_stat_once() { + let env = TestEnv::new(); + let mut ctx = env.ctx(); + let stats_backend = Arc::new(ControlledKvBackend::new()); + ctx.in_memory = stats_backend.clone(); + stats_backend.set_delay_puts(true); + + let flush_stats_factor = 3; + let handler = Arc::new(CollectStatsHandler::new(Some(flush_stats_factor))); + handle_stat(handler.clone(), ctx.clone(), stat(101, 1, 0, "dn-101")) + .await + .unwrap(); + stats_backend.clear_recorded_puts(); + + let mut tasks = Vec::with_capacity(2 * flush_stats_factor); + for marker in 1..=(2 * flush_stats_factor) { + tasks.push(tokio::spawn(handle_stat( + handler.clone(), + ctx.clone(), + stat(101, 1, marker as u64, "dn-101"), + ))); + } + for task in tasks { + timeout(Duration::from_secs(1), task) + .await + .unwrap() + .unwrap() + .unwrap(); + } + + let puts = stats_backend.recorded_puts(); + assert_eq!(2, puts.len()); + let mut markers = puts + .into_iter() + .flat_map(|put| { + let value: DatanodeStatValue = put.value.try_into().unwrap(); + value + .stats + .into_iter() + .map(|stat| stat.region_num) + .collect::>() + }) + .collect::>(); + markers.sort_unstable(); + assert_eq!((1..=6).collect::>(), markers); + assert_eq!(1, stats_backend.max_active_puts()); + } + async fn handle_request_many_times( mut ctx: Context, handler: &CollectStatsHandler,