diff --git a/rclrs/src/executor/basic_executor.rs b/rclrs/src/executor/basic_executor.rs index 4df0db9a1..4a4b42d55 100644 --- a/rclrs/src/executor/basic_executor.rs +++ b/rclrs/src/executor/basic_executor.rs @@ -136,20 +136,7 @@ impl ExecutorRuntime for BasicExecutorRuntime { })); while let Ok(task) = self.next_task(&workers_finished) { - // SAFETY: If the mutex is poisoned then we have unrecoverable situation. - let mut future_slot = task.future.lock().unwrap(); - if let Some(mut future) = future_slot.take() { - let waker = waker_ref(&task); - let task_context = &mut TaskContext::from_waker(&waker); - // Poll the future inside the task so it can do some work and - // tell us its state. - if future.as_mut().poll(task_context).is_pending() { - // The task is still pending, so return the future to its - // task so it can be processed again when it's ready to - // continue. - *future_slot = Some(future); - } - } + task.poll(); } let (runners, new_worker_receiver, errors) = worker_result_receiver.recv().expect( @@ -349,6 +336,7 @@ impl TaskSender { let task = Arc::new(Task { future: Mutex::new(Some(f)), task_sender: self.task_sender.clone(), + queued: AtomicBool::new(true), }); if let Err(_) = self.task_sender.send(task) { @@ -372,14 +360,44 @@ struct Task { /// shared borrow that comes from the Arc. future: Mutex>>, task_sender: Sender>, + /// Whether this task already has an entry in the ready queue. + queued: AtomicBool, } /// Implementing this trait gives us a very easy implementation of waking /// behavior for Task on our BasicExecutorRuntime. impl ArcWake for Task { fn wake_by_ref(arc_self: &Arc) { - let cloned = Arc::clone(arc_self); - arc_self.task_sender.send(cloned).ok(); + // Multiple wake-ups may be coalesced into one poll. Avoiding duplicate + // queue entries prevents a burst of wake-ups from causing redundant + // polls after the task has already observed the latest state. + if !arc_self.queued.swap(true, Ordering::AcqRel) { + if arc_self.task_sender.send(Arc::clone(arc_self)).is_err() { + arc_self.queued.store(false, Ordering::Release); + } + } + } +} + +impl Task { + fn poll(self: &Arc) { + // The queue entry has now been consumed, so a wake that happens while + // polling must schedule another poll. + self.queued.store(false, Ordering::Release); + + // SAFETY: If the mutex is poisoned then we have an unrecoverable situation. + let mut future_slot = self.future.lock().unwrap(); + if let Some(mut future) = future_slot.take() { + let waker = waker_ref(self); + let task_context = &mut TaskContext::from_waker(&waker); + // Poll the future inside the task so it can do some work and tell + // us its state. + if future.as_mut().poll(task_context).is_pending() { + // The task is still pending, so return the future to its task + // so it can be processed again when it's ready to continue. + *future_slot = Some(future); + } + } } } @@ -450,3 +468,68 @@ async fn manage_workers( (finished_runners, new_workers, errors) } + +#[cfg(test)] +mod tests { + use super::*; + use futures::future::poll_fn; + use std::{ + sync::atomic::AtomicUsize, + task::{Poll, Waker}, + }; + + #[test] + fn repeated_wakes_are_coalesced() { + let (task_sender, ready_queue) = channel(); + let task_sender = TaskSender { task_sender }; + let poll_count = Arc::new(AtomicUsize::new(0)); + let saved_waker = Arc::new(Mutex::new(None::)); + + let poll_count_clone = Arc::clone(&poll_count); + let saved_waker_clone = Arc::clone(&saved_waker); + task_sender.add_async_task(Box::pin(poll_fn(move |context| { + poll_count_clone.fetch_add(1, Ordering::Relaxed); + *saved_waker_clone.lock().unwrap() = Some(context.waker().clone()); + Poll::Pending + }))); + + ready_queue.try_recv().unwrap().poll(); + let waker = saved_waker.lock().unwrap().as_ref().unwrap().clone(); + for _ in 0..1_000 { + waker.wake_by_ref(); + } + + ready_queue.try_recv().unwrap().poll(); + assert_eq!(poll_count.load(Ordering::Relaxed), 2); + assert!( + ready_queue.try_recv().is_err(), + "A task should only have one outstanding queue entry" + ); + + // Break the cycle between the stored Waker and its Task. + saved_waker.lock().unwrap().take(); + } + + #[test] + fn wake_while_polling_schedules_another_poll() { + let (task_sender, ready_queue) = channel(); + let task_sender = TaskSender { task_sender }; + let poll_count = Arc::new(AtomicUsize::new(0)); + + let poll_count_clone = Arc::clone(&poll_count); + task_sender.add_async_task(Box::pin(poll_fn(move |context| { + if poll_count_clone.fetch_add(1, Ordering::Relaxed) == 0 { + context.waker().wake_by_ref(); + Poll::Pending + } else { + Poll::Ready(()) + } + }))); + + ready_queue.try_recv().unwrap().poll(); + ready_queue.try_recv().unwrap().poll(); + + assert_eq!(poll_count.load(Ordering::Relaxed), 2); + assert!(ready_queue.try_recv().is_err()); + } +}