Add method to drop TimerLoop

Author: foriequal0

util/timer/src/timer.rs

@@ -14,14 +14,15 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
+use std::cell::Cell;
 use std::cmp::Reverse;
 use std::collections::binary_heap::BinaryHeap;
 use std::collections::hash_map::{Entry, HashMap};
 use std::collections::VecDeque;
 use std::string::ToString;
 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::{Arc, Weak};
-use std::thread;
+use std::thread::{self, JoinHandle};
 use std::time::{Duration, Instant};
 
 use parking_lot::{Condvar, Mutex, RwLock};
@@ -36,28 +37,37 @@ pub trait TimeoutHandler: Send + Sync {
 #[derive(Clone)]
 pub struct TimerLoop {
     timer_id_nonce: Arc<AtomicUsize>,
+
     scheduler: Arc<Scheduler>,
+    scheduler_join_handle: Arc<Mutex<Cell<Option<JoinHandle<()>>>>>,
+
+    worker_queue: Arc<WorkerQueue>,
+    worker_join_handles: Arc<Mutex<Vec<JoinHandle<()>>>>,
 }
 
 impl TimerLoop {
     pub fn new(worker_size: usize) -> TimerLoop {
         let scheduler = Arc::new(Scheduler::new());
 
         let worker_queue = Arc::new(WorkerQueue::new());
-
-        spawn_workers(worker_size, &worker_queue);
-        {
+        let worker_join_handles = spawn_workers(worker_size, &worker_queue);
+        let scheduler_join_handle = {
             let worker_queue = Arc::clone(&worker_queue);
             let scheduler = Arc::clone(&scheduler);
             thread::Builder::new()
                 .name("timer.scheduler".to_string())
                 .spawn(move || scheduler.run(&worker_queue))
-                .unwrap();
-        }
+                .unwrap()
+        };
 
         TimerLoop {
             timer_id_nonce: Arc::new(AtomicUsize::new(0)),
+
             scheduler,
+            scheduler_join_handle: Arc::new(Mutex::new(Cell::new(Some(scheduler_join_handle)))),
+
+            worker_queue,
+            worker_join_handles: Arc::new(Mutex::new(worker_join_handles)),
         }
     }
 
@@ -70,6 +80,19 @@ impl TimerLoop {
             scheduler: Arc::downgrade(&self.scheduler),
         }
     }
+
+    pub fn stop(self) {
+        self.scheduler.notify_stop();
+        self.worker_queue.notify_stop();
+
+        if let Some(join_handle) = self.scheduler_join_handle.lock().take() {
+            let _ = join_handle.join();
+        }
+        let mut worker_join_handles = self.worker_join_handles.lock();
+        for join_handle in worker_join_handles.drain(..) {
+            let _ = join_handle.join();
+        }
+    }
 }
 
 type TimerId = usize;
@@ -147,6 +170,12 @@ impl Scheduler {
         result
     }
 
+    fn notify_stop(&self) {
+        let mut scheduler = self.inner.lock();
+        scheduler.stop();
+        self.condvar.notify_all();
+    }
+
     fn run(&self, worker_queue: &WorkerQueue) {
         let mut scheduler = self.inner.lock();
         while !scheduler.stop {
@@ -158,6 +187,7 @@ impl Scheduler {
                 None => self.condvar.wait(&mut scheduler),
             }
         }
+        ctrace!(TIMER, "Scheduler loop has been stopped");
     }
 }
 
@@ -304,6 +334,14 @@ impl SchedulerInner {
         }
     }
 
+    fn stop(&mut self) {
+        self.stop = true;
+        self.states.clear();
+        for Reverse(TimeOrdered(schedule)) in self.heap.drain() {
+            schedule.state_control.cancel();
+        }
+    }
+
     fn handle_timeout(&mut self, worker_queue: &WorkerQueue) -> Option<Duration> {
         loop {
             let now = Instant::now();
@@ -531,11 +569,15 @@ impl Callback {
     }
 }
 
-fn spawn_workers(size: usize, queue: &Arc<WorkerQueue>) {
+fn spawn_workers(size: usize, queue: &Arc<WorkerQueue>) -> Vec<JoinHandle<()>> {
+    let mut handles = Vec::new();
     for i in 0..size {
         let queue = Arc::clone(queue);
-        thread::Builder::new().name(format!("timer.worker.{}", i)).spawn(move || worker_loop(&queue)).unwrap();
+        let handle =
+            thread::Builder::new().name(format!("timer.worker.{}", i)).spawn(move || worker_loop(&queue)).unwrap();
+        handles.push(handle);
     }
+    handles
 }
 
 fn worker_loop(queue: &Arc<WorkerQueue>) {
@@ -565,6 +607,7 @@ fn worker_loop(queue: &Arc<WorkerQueue>) {
             }
         });
     }
+    ctrace!(TIMER, "Worker loop has been stopped");
 }
 
 struct WorkerQueue {
@@ -601,6 +644,13 @@ impl WorkerQueue {
         }
         queue.pop_front()
     }
+
+    fn notify_stop(&self) {
+        let mut queue = self.queue.lock();
+        queue.clear();
+        self.stop.store(true, Ordering::SeqCst);
+        self.condvar.notify_all();
+    }
 }
 
 #[cfg(test)]
@@ -676,6 +726,7 @@ mod tests {
         let (called_at, token) = value.unwrap();
         assert_eq!(token, timer_token);
         assert!(similar(called_at, begin + tick())); // called_at = now + ticksufficiently small
+        timer_loop.stop();
     }
 
     #[test]
@@ -701,6 +752,7 @@ mod tests {
         let mut value = mutex.lock();
         condvar.wait_for(&mut value, long_tick());
         assert!(value.is_none());
+        timer_loop.stop();
     }
 
     #[test]
@@ -747,6 +799,7 @@ mod tests {
 
         assert_eq!(timer.schedule_once(tick(), timer_token), Ok(()));
         assert_eq!(timer.schedule_once(tick(), timer_token), Err(ScheduleError::TokenAlreadyScheduled));
+        timer_loop.stop();
     }
 
     #[test]
@@ -759,6 +812,7 @@ mod tests {
 
         assert_eq!(timer.schedule_once(tick(), timer_token_1), Ok(()));
         assert_eq!(timer.schedule_once(tick(), timer_token_2), Ok(()));
+        timer_loop.stop();
     }
 
     #[test]
@@ -790,6 +844,7 @@ mod tests {
         let (called_at, token) = value.unwrap();
         assert_eq!(token, timer_token);
         assert!(similar(called_at, begin + tick())); // called_at = now + ticksufficiently small
+        timer_loop.stop();
     }
 
     #[test]
@@ -821,6 +876,7 @@ mod tests {
         let (called_at, token) = value.unwrap();
         assert_eq!(token, timer_token);
         assert!(similar(called_at, begin + tick())); // called_at = now + ticksufficiently small
+        timer_loop.stop();
     }
     #[test]
     fn test_repeat() {
@@ -860,5 +916,31 @@ mod tests {
         for i in 1..TEST_COUNT {
             assert!(similar(value[i - 1] + tick(), value[i]));
         }
+        timer_loop.stop();
+    }
+
+    #[test]
+    fn test_timerloop_drop() {
+        let timer_token = 100;
+        let timer_loop = Box::new(TimerLoop::new(4));
+        let pair = Arc::new((Condvar::new(), Mutex::new(None)));
+        let handler = {
+            let pair = Arc::clone(&pair);
+            Arc::new(CallbackHandler(move |_| {
+                let (ref condvar, ref mutex) = *pair;
+                let mut value = mutex.lock();
+                *value = Some(());
+                condvar.notify_all();
+            }))
+        };
+        let timer = new_timer(&timer_loop, "test", &handler);
+
+        assert_eq!(timer.schedule_once(tick(), timer_token), Ok(()));
+        timer_loop.stop();
+
+        let (ref condvar, ref mutex) = *pair;
+        let mut value = mutex.lock();
+        condvar.wait_for(&mut value, long_tick());
+        assert!(value.is_none());
     }
 }