Tiny extract PeriodicTask in router (#15988)

sgl-model-gateway/src/policies/cache_aware.rs

@@ -59,22 +59,15 @@
     during the next eviction cycle.
 */
 
-use std::{
+use std::sync::Arc;
-    sync::{
-        atomic::{AtomicBool, Ordering},
-        Arc,
-    },
-    thread,
-    time::Duration,
-};
 
 use dashmap::DashMap;
 use rand::Rng;
 use tracing::debug;
 
 use super::{
-    get_healthy_worker_indices, normalize_model_key, tree::Tree, CacheAwareConfig,
+    get_healthy_worker_indices, normalize_model_key, tree::Tree, utils::PeriodicTask,
-    LoadBalancingPolicy, SelectWorkerInfo,
+    CacheAwareConfig, LoadBalancingPolicy, SelectWorkerInfo,
 };
 use crate::core::Worker;
 
@@ -87,10 +80,7 @@ use crate::core::Worker;
 pub struct CacheAwarePolicy {
     config: CacheAwareConfig,
     trees: Arc<DashMap<String, Arc<Tree>>>,
-    /// Handle to the background eviction thread
+    _eviction_task: Option<PeriodicTask>,
-    eviction_handle: Option<thread::JoinHandle<()>>,
-    /// Flag to signal the eviction thread to stop
-    shutdown_flag: Arc<AtomicBool>,
 }
 
 impl CacheAwarePolicy {
@@ -100,39 +90,16 @@ impl CacheAwarePolicy {
 
     pub fn with_config(config: CacheAwareConfig) -> Self {
         let trees = Arc::new(DashMap::<String, Arc<Tree>>::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 || {
+            Some(PeriodicTask::spawn(
-                // Use smaller sleep intervals to check shutdown flag more frequently
+                config.eviction_interval_secs,
-                let check_interval_ms = 100; // Check every 100ms
+                "Eviction",
-                let total_sleep_ms = interval * 1000;
+                move || {
-
-                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
                     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,
+            _eviction_task: eviction_task,
-            shutdown_flag,
         }
     }
 
@@ -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::*;

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;

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<AtomicBool>,
+    handle: Option<JoinHandle<()>>,
+}
+
+impl PeriodicTask {
+    /// Spawn a background thread that periodically executes a task.
+    pub fn spawn<F>(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));
+    }
+}