Merge pull request #8 from de-vri-es/fix-waker-list-growing-indefinitely

Fix the waker lists growing indefinitely.

Author: de-vri-es

CHANGELOG.md

@@ -1,5 +1,8 @@
+# main
+* Fix bug where the list of wakers to trigger on shutdown or shutdown completion could grow indefinitely.
+
 # Version 0.2.1 - 2023-10-08
-* Fix `ShutdownManager::wait_shutdown_complete()` never completing if callend when no shutdown was triggered yet and no delay tokens exist.
+* Fix `ShutdownManager::wait_shutdown_complete()` never completing if called when no shutdown was triggered yet and no delay tokens exist.
 
 # Version 0.2.0 - 2023-09-26:
 * Rename `Shutdown` struct to `ShutdownManager`.

src/lib.rs

@@ -169,7 +169,6 @@
 
 use std::future::Future;
 use std::sync::{Arc, Mutex};
-use std::task::Waker;
 
 mod shutdown_complete;
 pub use shutdown_complete::ShutdownComplete;
@@ -178,6 +177,7 @@ mod shutdown_signal;
 pub use shutdown_signal::ShutdownSignal;
 
 mod wrap_cancel;
+use waker_list::WakerList;
 pub use wrap_cancel::WrapCancel;
 
 mod wrap_trigger_shutdown;
@@ -186,6 +186,8 @@ pub use wrap_trigger_shutdown::WrapTriggerShutdown;
 mod wrap_delay_shutdown;
 pub use wrap_delay_shutdown::WrapDelayShutdown;
 
+mod waker_list;
+
 /// Shutdown manager for asynchronous tasks and futures.
 ///
 /// The shutdown manager allows you to:
@@ -244,6 +246,7 @@ impl<T: Clone> ShutdownManager<T> {
 	pub fn wait_shutdown_triggered(&self) -> ShutdownSignal<T> {
 		ShutdownSignal {
 			inner: self.inner.clone(),
+			waker_token: None,
 		}
 	}
 
@@ -258,6 +261,7 @@ impl<T: Clone> ShutdownManager<T> {
 	pub fn wait_shutdown_complete(&self) -> ShutdownComplete<T> {
 		ShutdownComplete {
 			inner: self.inner.clone(),
+			waker_token: None,
 		}
 	}
 
@@ -459,19 +463,19 @@ struct ShutdownManagerInner<T> {
 	delay_tokens: usize,
 
 	/// Tasks to wake when a shutdown is triggered.
-	on_shutdown: Vec<Waker>,
+	on_shutdown: WakerList,
 
 	/// Tasks to wake when the shutdown is complete.
-	on_shutdown_complete: Vec<Waker>,
+	on_shutdown_complete: WakerList,
 }
 
 impl<T: Clone> ShutdownManagerInner<T> {
 	fn new() -> Self {
 		Self {
 			shutdown_reason: None,
 			delay_tokens: 0,
-			on_shutdown_complete: Vec::new(),
-			on_shutdown: Vec::new(),
+			on_shutdown_complete: WakerList::new(),
+			on_shutdown: WakerList::new(),
 		}
 	}
 
@@ -493,9 +497,7 @@ impl<T: Clone> ShutdownManagerInner<T> {
 			},
 			None => {
 				self.shutdown_reason = Some(reason);
-				for abort in std::mem::take(&mut self.on_shutdown) {
-					abort.wake()
-				}
+				self.on_shutdown.wake_all();
 				if self.delay_tokens == 0 {
 					self.notify_shutdown_complete()
 				}
@@ -505,9 +507,7 @@ impl<T: Clone> ShutdownManagerInner<T> {
 	}
 
 	fn notify_shutdown_complete(&mut self) {
-		for waiter in std::mem::take(&mut self.on_shutdown_complete) {
-			waiter.wake()
-		}
+		self.on_shutdown_complete.wake_all();
 	}
 }
 

src/shutdown_complete.rs

@@ -3,30 +3,137 @@ use std::pin::Pin;
 use std::sync::{Arc, Mutex};
 use std::task::{Context, Poll};
 
+use crate::waker_list::WakerToken;
 use crate::ShutdownManagerInner;
 
 /// Future to wait for a shutdown to complete.
 pub struct ShutdownComplete<T: Clone> {
 	pub(crate) inner: Arc<Mutex<ShutdownManagerInner<T>>>,
+	pub(crate) waker_token: Option<WakerToken>,
+}
+
+impl<T: Clone> Clone for ShutdownComplete<T> {
+	fn clone(&self) -> Self {
+		// Clone only the reference to the shutdown manager, not the waker token.
+		// The waker token is personal to each future.
+		Self {
+			inner: self.inner.clone(),
+			waker_token: None,
+		}
+	}
+}
+
+impl<T: Clone> Drop for ShutdownComplete<T> {
+	fn drop(&mut self) {
+		if let Some(token) = self.waker_token.take() {
+			let mut inner = self.inner.lock().unwrap();
+			inner.on_shutdown_complete.deregister(token);
+		}
+	}
 }
 
 impl<T: Clone> Future for ShutdownComplete<T> {
 	type Output = T;
 
 	#[inline]
 	fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll<Self::Output> {
-		let me = self.as_ref();
+		let me = self.get_mut();
 		let mut inner = me.inner.lock().unwrap();
+
+		// We're being polled, so we should deregister the waker (if any).
+		if let Some(token) = me.waker_token.take() {
+			inner.on_shutdown_complete.deregister(token);
+		}
+
+		// Check if the shutdown is completed.
 		if inner.delay_tokens == 0 {
 			if let Some(reason) = inner.shutdown_reason.clone() {
-				Poll::Ready(reason)
-			} else {
-				inner.on_shutdown_complete.push(context.waker().clone());
-				Poll::Pending
+				return Poll::Ready(reason);
 			}
-		} else {
-			inner.on_shutdown_complete.push(context.waker().clone());
-			Poll::Pending
+		}
+
+		// We're not ready, so register the waker to wake us on shutdown completion.
+		me.waker_token = Some(inner.on_shutdown_complete.register(context.waker().clone()));
+
+		Poll::Pending
+	}
+}
+
+#[cfg(test)]
+mod test {
+	use assert2::assert;
+	use std::future::Future;
+	use std::pin::Pin;
+	use std::task::Poll;
+
+	/// Wrapper around a future to poll it only once.
+	struct PollOnce<'a, F>(&'a mut F);
+
+	impl<'a, F: std::marker::Unpin + Future> Future for PollOnce<'a, F> {
+		type Output = Poll<F::Output>;
+
+		fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
+			Poll::Ready(Pin::new(&mut self.get_mut().0).poll(cx))
+		}
+	}
+
+	/// Poll a future once.
+	async fn poll_once<F: Future + Unpin>(future: &mut F) -> Poll<F::Output> {
+		PollOnce(future).await
+	}
+
+	#[tokio::test]
+	async fn waker_list_doesnt_grow_infinitely() {
+		let shutdown = crate::ShutdownManager::<()>::new();
+		for i in 0..100_000 {
+			let mut wait_shutdown_complete = shutdown.wait_shutdown_complete();
+			let task = tokio::spawn(async move {
+				assert!(let Poll::Pending = poll_once(&mut wait_shutdown_complete).await);
+			});
+			assert!(let Ok(()) = task.await, "task = {i}");
+		}
+
+		// Since we wait for each task to complete before spawning another,
+		// the total amount of waker slots used should be only 1.
+		let inner = shutdown.inner.lock().unwrap();
+		assert!(inner.on_shutdown_complete.total_slots() == 1);
+		assert!(inner.on_shutdown_complete.empty_slots() == 1);
+	}
+
+	#[tokio::test]
+	async fn cloning_does_not_clone_waker_token() {
+		let shutdown = crate::ShutdownManager::<()>::new();
+
+		let mut signal = shutdown.wait_shutdown_complete();
+		assert!(let None = &signal.waker_token);
+
+		assert!(let Poll::Pending = poll_once(&mut signal).await);
+		assert!(let Some(_) = &signal.waker_token);
+
+		let mut cloned = signal.clone();
+		assert!(let None = &cloned.waker_token);
+		assert!(let Some(_) = &signal.waker_token);
+
+		assert!(let Poll::Pending = poll_once(&mut cloned).await);
+		assert!(let Some(_) = &cloned.waker_token);
+		assert!(let Some(_) = &signal.waker_token);
+
+		{
+			let inner = shutdown.inner.lock().unwrap();
+			assert!(inner.on_shutdown_complete.total_slots() == 2);
+			assert!(inner.on_shutdown_complete.empty_slots() == 0);
+		}
+
+		{
+			drop(signal);
+			let inner = shutdown.inner.lock().unwrap();
+			assert!(inner.on_shutdown_complete.empty_slots() == 1);
+		}
+
+		{
+			drop(cloned);
+			let inner = shutdown.inner.lock().unwrap();
+			assert!(inner.on_shutdown_complete.empty_slots() == 2);
 		}
 	}
 }

src/shutdown_signal.rs

@@ -3,16 +3,37 @@ use std::pin::Pin;
 use std::sync::{Arc, Mutex};
 use std::task::{Context, Poll};
 
+use crate::waker_list::WakerToken;
 use crate::{WrapCancel, ShutdownManagerInner};
 
 /// A future to wait for a shutdown signal.
 ///
 /// The future completes when the associated [`ShutdownManager`][crate::ShutdownManager] triggers a shutdown.
 ///
 /// The shutdown signal can be cloned and sent between threads freely.
-#[derive(Clone)]
 pub struct ShutdownSignal<T: Clone> {
 	pub(crate) inner: Arc<Mutex<ShutdownManagerInner<T>>>,
+	pub(crate) waker_token: Option<WakerToken>,
+}
+
+impl<T: Clone> Clone for ShutdownSignal<T> {
+	fn clone(&self) -> Self {
+		// Clone only the reference to the shutdown manager, not the waker token.
+		// The waker token is personal to each future.
+		Self {
+			inner: self.inner.clone(),
+			waker_token: None,
+		}
+	}
+}
+
+impl<T: Clone> Drop for ShutdownSignal<T> {
+	fn drop(&mut self) {
+		if let Some(token) = self.waker_token.take() {
+			let mut inner = self.inner.lock().unwrap();
+			inner.on_shutdown.deregister(token);
+		}
+	}
 }
 
 impl<T: Clone> ShutdownSignal<T> {
@@ -36,13 +57,99 @@ impl<T: Clone> Future for ShutdownSignal<T> {
 
 	#[inline]
 	fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll<Self::Output> {
-		let me = self.as_ref();
+		let me = self.get_mut();
 		let mut inner = me.inner.lock().unwrap();
+
+		// We're being polled, so we should deregister the waker (if any).
+		if let Some(token) = me.waker_token.take() {
+			inner.on_shutdown.deregister(token);
+		}
+
 		if let Some(reason) = inner.shutdown_reason.clone() {
+			// Shutdown started, so we're ready.
 			Poll::Ready(reason)
 		} else {
-			inner.on_shutdown.push(context.waker().clone());
+			// We're not ready, so register the waker to wake us on shutdown start.
+			me.waker_token = Some(inner.on_shutdown.register(context.waker().clone()));
 			Poll::Pending
 		}
 	}
 }
+
+#[cfg(test)]
+mod test {
+	use assert2::assert;
+	use std::future::Future;
+	use std::pin::Pin;
+	use std::task::Poll;
+
+	/// Wrapper around a future to poll it only once.
+	struct PollOnce<'a, F>(&'a mut F);
+
+	impl<'a, F: std::marker::Unpin + Future> Future for PollOnce<'a, F> {
+		type Output = Poll<F::Output>;
+
+		fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
+			Poll::Ready(Pin::new(&mut self.get_mut().0).poll(cx))
+		}
+	}
+
+	/// Poll a future once.
+	async fn poll_once<F: Future + Unpin>(future: &mut F) -> Poll<F::Output> {
+		PollOnce(future).await
+	}
+
+	#[tokio::test]
+	async fn waker_list_doesnt_grow_infinitely() {
+		let shutdown = crate::ShutdownManager::<()>::new();
+		for i in 0..100_000 {
+			let task = tokio::spawn(shutdown.wrap_cancel(async move {
+				tokio::task::yield_now().await;
+			}));
+			assert!(let Ok(Ok(())) = task.await, "task = {i}");
+		}
+
+		// Since we wait for each task to complete before spawning another,
+		// the total amount of waker slots used should be only 1.
+		let inner = shutdown.inner.lock().unwrap();
+		assert!(inner.on_shutdown.total_slots() == 1);
+		assert!(inner.on_shutdown.empty_slots() == 1);
+	}
+
+	#[tokio::test]
+	async fn cloning_does_not_clone_waker_token() {
+		let shutdown = crate::ShutdownManager::<()>::new();
+
+		let mut signal = shutdown.wait_shutdown_triggered();
+		assert!(let None = &signal.waker_token);
+
+		assert!(let Poll::Pending = poll_once(&mut signal).await);
+		assert!(let Some(_) = &signal.waker_token);
+
+		let mut cloned = signal.clone();
+		assert!(let None = &cloned.waker_token);
+		assert!(let Some(_) = &signal.waker_token);
+
+		assert!(let Poll::Pending = poll_once(&mut cloned).await);
+		assert!(let Some(_) = &cloned.waker_token);
+		assert!(let Some(_) = &signal.waker_token);
+
+		{
+			let inner = shutdown.inner.lock().unwrap();
+			assert!(inner.on_shutdown.total_slots() == 2);
+			assert!(inner.on_shutdown.empty_slots() == 0);
+		}
+
+		{
+			drop(signal);
+			let inner = shutdown.inner.lock().unwrap();
+			assert!(inner.on_shutdown.empty_slots() == 1);
+		}
+
+		{
+			drop(cloned);
+			let inner = shutdown.inner.lock().unwrap();
+			assert!(inner.on_shutdown.empty_slots() == 2);
+		}
+	}
+}