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};
@@ -38,28 +39,35 @@ pub struct TimerLoop {
     timer_id_nonce: AtomicUsize,
 
     scheduler: Arc<Scheduler>,
+    scheduler_join_handle: Mutex<Cell<Option<JoinHandle<()>>>>,
+
+    worker_queue: Arc<WorkerQueue>,
+    worker_join_handles: 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: AtomicUsize::new(0),
 
             scheduler,
+            scheduler_join_handle: Mutex::new(Cell::new(Some(scheduler_join_handle))),
+
+            worker_queue,
+            worker_join_handles: Mutex::new(worker_join_handles),
         }
     }
 
@@ -78,6 +86,30 @@ impl TimerLoop {
         timer.set_name(name);
         timer
     }
+
+    /* stop() will clear all remaining schedules,
+     * while drop() tries to deliver remaining schedules */
+    pub fn stop(self) {
+        self.scheduler.clear_and_stop();
+        self.worker_queue.clear_and_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();
+        }
+    }
+}
+
+impl Drop for TimerLoop {
+    fn drop(&mut self) {
+        self.scheduler.start_drain();
+        // worker will be handled after the scheduler finishes its job.
+        // The scheduler thread will be the only strong ref holder after this.
+        // (all TimerApi will holds weak refs only)
+    }
 }
 
 type TimerId = usize;
@@ -161,17 +193,37 @@ impl Scheduler {
         result
     }
 
+    fn clear_and_stop(&self) {
+        let mut scheduler = self.inner.lock();
+        scheduler.clear_and_stop();
+        self.condvar.notify_all();
+    }
+
+    fn start_drain(&self) {
+        let mut scheduler = self.inner.lock();
+        scheduler.finished = true;
+        self.condvar.notify_all();
+    }
+
     fn run(&self, worker_queue: &WorkerQueue) {
         let mut scheduler = self.inner.lock();
-        while !scheduler.stop {
+        loop {
             let wait_for = scheduler.handle_timeout(worker_queue);
             match wait_for {
                 Some(timeout) => {
                     self.condvar.wait_for(&mut scheduler, timeout);
                 }
-                None => self.condvar.wait(&mut scheduler),
+                None => {
+                    if scheduler.finished {
+                        break
+                    } else {
+                        self.condvar.wait(&mut scheduler)
+                    }
+                }
             }
         }
+        worker_queue.enqueue_finished();
+        ctrace!(TIMER, "Scheduler loop has been stopped");
     }
 }
 
@@ -208,14 +260,15 @@ impl Scheduler {
 struct SchedulerInner {
     states: HashMap<ScheduleId, Arc<ScheduleStateControl>>,
     heap: BinaryHeap<Reverse<TimeOrdered<Schedule>>>,
-    stop: bool,
+    finished: bool,
 }
+
 impl SchedulerInner {
     fn new() -> SchedulerInner {
         SchedulerInner {
             states: HashMap::new(),
             heap: BinaryHeap::new(),
-            stop: false,
+            finished: false,
         }
     }
 
@@ -226,6 +279,11 @@ impl SchedulerInner {
         after: Duration,
         repeat: Option<Duration>,
     ) -> Result<(), ScheduleError> {
+        if self.finished {
+            cdebug!(TIMER, "schedule: TimerLoop has been finished");
+            return Ok(())
+        }
+
         let schedule_id = ScheduleId(requested_timer.timer_id, timer_token);
         let handler = {
             let guard = requested_timer.handler.read();
@@ -320,6 +378,14 @@ impl SchedulerInner {
         }
     }
 
+    fn clear_and_stop(&mut self) {
+        self.finished = 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();
@@ -547,11 +613,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::with_capacity(size);
     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>) {
@@ -581,26 +651,27 @@ fn worker_loop(queue: &Arc<WorkerQueue>) {
             }
         });
     }
+    ctrace!(TIMER, "Worker loop has been stopped");
 }
 
 struct WorkerQueue {
     queue: Mutex<VecDeque<Callback>>,
     condvar: Condvar,
-    stop: AtomicBool,
+    finished: AtomicBool,
 }
 
 impl WorkerQueue {
     fn new() -> WorkerQueue {
         WorkerQueue {
             queue: Mutex::new(VecDeque::new()),
             condvar: Condvar::new(),
-            stop: AtomicBool::new(false),
+            finished: AtomicBool::new(false),
         }
     }
 
     fn enqueue(&self, callback: Callback) {
         let mut queue = self.queue.lock();
-        if self.stop.load(Ordering::SeqCst) {
+        if self.finished.load(Ordering::SeqCst) {
             return
         }
         queue.push_back(callback);
@@ -610,13 +681,25 @@ impl WorkerQueue {
     fn dequeue(&self) -> Option<Callback> {
         let mut queue = self.queue.lock();
         while queue.is_empty() {
-            if self.stop.load(Ordering::SeqCst) {
+            if self.finished.load(Ordering::SeqCst) {
                 return None
             }
             self.condvar.wait(&mut queue);
         }
         queue.pop_front()
     }
+
+    fn clear_and_stop(&self) {
+        let mut queue = self.queue.lock();
+        queue.clear();
+        self.finished.store(true, Ordering::SeqCst);
+        self.condvar.notify_all();
+    }
+
+    fn enqueue_finished(&self) {
+        self.finished.store(true, Ordering::SeqCst);
+        self.condvar.notify_all();
+    }
 }
 
 #[cfg(test)]
@@ -693,6 +776,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]
@@ -718,6 +802,7 @@ mod tests {
         let mut value = mutex.lock();
         condvar.wait_for(&mut value, long_tick());
         assert!(value.is_none());
+        timer_loop.stop();
     }
 
     #[test]
@@ -773,6 +858,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]
@@ -785,6 +871,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]
@@ -816,6 +903,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]
@@ -847,6 +935,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() {
@@ -886,5 +975,59 @@ mod tests {
         for i in 1..TEST_COUNT {
             assert!(similar(value[i - 1] + tick(), value[i]));
         }
+        timer_loop.stop();
+    }
+
+    #[test]
+    fn test_timerloop_stop() {
+        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());
+    }
+
+
+    #[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(()));
+        drop(timer_loop);
+
+        let (ref condvar, ref mutex) = *pair;
+        let mut value = mutex.lock();
+        condvar.wait_for(&mut value, long_tick());
+        assert_eq!(*value, Some(()));
+
+        assert_eq!(timer.schedule_once(tick(), timer_token), Err(ScheduleError::TimerLoopDropped));
     }
 }