From 26c50912175ee7d2e94adaedefa6bb98459a51f2 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 28 Dec 2025 14:07:13 +0800 Subject: [PATCH] Tiny extract PeriodicTask in router (#15988) --- sgl-model-gateway/src/policies/cache_aware.rs | 74 ++---------- sgl-model-gateway/src/policies/mod.rs | 1 + sgl-model-gateway/src/policies/utils.rs | 110 ++++++++++++++++++ 3 files changed, 123 insertions(+), 62 deletions(-) create mode 100644 sgl-model-gateway/src/policies/utils.rs diff --git a/sgl-model-gateway/src/policies/cache_aware.rs b/sgl-model-gateway/src/policies/cache_aware.rs index 73d75f35c..e3c5239f3 100644 --- a/sgl-model-gateway/src/policies/cache_aware.rs +++ b/sgl-model-gateway/src/policies/cache_aware.rs @@ -59,22 +59,15 @@ during the next eviction cycle. */ -use std::{ - sync::{ - atomic::{AtomicBool, Ordering}, - Arc, - }, - thread, - time::Duration, -}; +use std::sync::Arc; use dashmap::DashMap; use rand::Rng; use tracing::debug; use super::{ - get_healthy_worker_indices, normalize_model_key, tree::Tree, CacheAwareConfig, - LoadBalancingPolicy, SelectWorkerInfo, + get_healthy_worker_indices, normalize_model_key, tree::Tree, utils::PeriodicTask, + CacheAwareConfig, LoadBalancingPolicy, SelectWorkerInfo, }; use crate::core::Worker; @@ -87,10 +80,7 @@ use crate::core::Worker; pub struct CacheAwarePolicy { config: CacheAwareConfig, trees: Arc>>, - /// Handle to the background eviction thread - eviction_handle: Option>, - /// Flag to signal the eviction thread to stop - shutdown_flag: Arc, + _eviction_task: Option, } impl CacheAwarePolicy { @@ -100,39 +90,16 @@ impl CacheAwarePolicy { pub fn with_config(config: CacheAwareConfig) -> Self { let trees = Arc::new(DashMap::>::new()); - let shutdown_flag = Arc::new(AtomicBool::new(false)); // Start background eviction thread if configured - let eviction_handle = if config.eviction_interval_secs > 0 { + let eviction_task = if config.eviction_interval_secs > 0 { let trees_clone = Arc::clone(&trees); - let shutdown_clone = Arc::clone(&shutdown_flag); let max_tree_size = config.max_tree_size; - let interval = config.eviction_interval_secs; - Some(thread::spawn(move || { - // Use smaller sleep intervals to check shutdown flag more frequently - let check_interval_ms = 100; // Check every 100ms - let total_sleep_ms = interval * 1000; - - loop { - // Sleep in small increments, checking shutdown flag periodically - let mut slept_ms = 0u64; - while slept_ms < total_sleep_ms { - if shutdown_clone.load(Ordering::Relaxed) { - debug!("Eviction thread received shutdown signal"); - return; - } - thread::sleep(Duration::from_millis(check_interval_ms)); - slept_ms += check_interval_ms; - } - - // Check shutdown before starting eviction - if shutdown_clone.load(Ordering::Relaxed) { - debug!("Eviction thread received shutdown signal"); - return; - } - - // Evict for all model trees + Some(PeriodicTask::spawn( + config.eviction_interval_secs, + "Eviction", + move || { for tree_ref in trees_clone.iter() { let model_id = tree_ref.key(); let tree = tree_ref.value(); @@ -143,8 +110,8 @@ impl CacheAwarePolicy { model_id, max_tree_size ); } - } - })) + }, + )) } else { None }; @@ -152,8 +119,7 @@ impl CacheAwarePolicy { Self { config, trees, - eviction_handle, - shutdown_flag, + _eviction_task: eviction_task, } } @@ -408,22 +374,6 @@ impl Default for CacheAwarePolicy { } } -impl Drop for CacheAwarePolicy { - fn drop(&mut self) { - // Signal the eviction thread to stop - self.shutdown_flag.store(true, Ordering::Relaxed); - - // Wait for the thread to finish (with timeout) - if let Some(handle) = self.eviction_handle.take() { - // The thread checks the shutdown flag every 100ms, so it should exit quickly - match handle.join() { - Ok(()) => debug!("Eviction thread shut down cleanly"), - Err(_) => debug!("Eviction thread panicked during shutdown"), - } - } - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/sgl-model-gateway/src/policies/mod.rs b/sgl-model-gateway/src/policies/mod.rs index 709decf71..3ad4b47ae 100644 --- a/sgl-model-gateway/src/policies/mod.rs +++ b/sgl-model-gateway/src/policies/mod.rs @@ -18,6 +18,7 @@ mod random; mod registry; mod round_robin; pub mod tree; +pub(crate) mod utils; pub use bucket::BucketPolicy; pub use cache_aware::CacheAwarePolicy; pub use consistent_hashing::ConsistentHashingPolicy; diff --git a/sgl-model-gateway/src/policies/utils.rs b/sgl-model-gateway/src/policies/utils.rs new file mode 100644 index 000000000..3e2c3b0a1 --- /dev/null +++ b/sgl-model-gateway/src/policies/utils.rs @@ -0,0 +1,110 @@ +use std::{ + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + thread::{self, JoinHandle}, + time::Duration, +}; + +use tracing::debug; + +#[derive(Debug)] +pub struct PeriodicTask { + debug_name: &'static str, + shutdown_flag: Arc, + handle: Option>, +} + +impl PeriodicTask { + /// Spawn a background thread that periodically executes a task. + pub fn spawn(interval_secs: u64, debug_name: &'static str, task: F) -> Self + where + F: Fn() + Send + 'static, + { + let shutdown_flag = Arc::new(AtomicBool::new(false)); + let shutdown_clone = Arc::clone(&shutdown_flag); + + let handle = thread::spawn(move || { + let check_interval_ms = 100u64; + let total_sleep_ms = interval_secs * 1000; + + loop { + // Sleep in small increments, checking shutdown flag periodically + let mut slept_ms = 0u64; + while slept_ms < total_sleep_ms { + if shutdown_clone.load(Ordering::Relaxed) { + debug!("{} thread received shutdown signal", debug_name); + return; + } + thread::sleep(Duration::from_millis(check_interval_ms)); + slept_ms += check_interval_ms; + } + + // Check shutdown before starting task + if shutdown_clone.load(Ordering::Relaxed) { + debug!("{} thread received shutdown signal", debug_name); + return; + } + + task(); + } + }); + + Self { + debug_name, + shutdown_flag, + handle: Some(handle), + } + } +} + +impl Drop for PeriodicTask { + fn drop(&mut self) { + self.shutdown_flag.store(true, Ordering::Relaxed); + + if let Some(handle) = self.handle.take() { + match handle.join() { + Ok(()) => debug!("{} thread shut down cleanly", self.debug_name), + Err(_) => debug!("{} thread panicked during shutdown", self.debug_name), + } + } + } +} + +#[cfg(test)] +mod tests { + use std::{sync::atomic::AtomicUsize, time::Instant}; + + use super::*; + + #[test] + fn test_periodic_task_executes() { + let counter = Arc::new(AtomicUsize::new(0)); + let counter_clone = Arc::clone(&counter); + + let _task = PeriodicTask::spawn(1, "test", move || { + counter_clone.fetch_add(1, Ordering::SeqCst); + }); + + // Wait for at least one execution + thread::sleep(Duration::from_millis(1200)); + assert!(counter.load(Ordering::SeqCst) >= 1); + + // Task will be stopped on drop + } + + #[test] + fn test_periodic_task_responds_to_shutdown() { + let task = PeriodicTask::spawn(60, "test", || { + // Long interval task + }); + + let start = Instant::now(); + drop(task); + let elapsed = start.elapsed(); + + // Should shutdown within ~200ms (2 check intervals), not 60 seconds + assert!(elapsed < Duration::from_millis(500)); + } +}