Make the notify() function return notified count

Fix off-by-one error in implementation
smol-rs · May 17, 2023 · 6345794 · 6345794
1 parent fb90f08
commit 6345794
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 28 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -360,7 +360,7 @@ impl<T> Event<T> {
     /// event.notify(1.additional().relaxed());
     /// ```
     #[inline]
-    pub fn notify(&self, notify: impl IntoNotification<Tag = T>) {
+    pub fn notify(&self, notify: impl IntoNotification<Tag = T>) -> usize {
         let notify = notify.into_notification();
 
         // Make sure the notification comes after whatever triggered it.
@@ -376,9 +376,11 @@ impl<T> Event<T> {
             // Notify if there is at least one unnotified listener and the number of notified
             // listeners is less than `limit`.
             if inner.notified.load(Ordering::Acquire) < limit {
-                inner.notify(notify);
+                return inner.notify(notify);
             }
         }
+
+        0
     }
 
     /// Return a reference to the inner state if it has been initialized.
@@ -489,7 +491,7 @@ impl Event<()> {
     /// event.notify_relaxed(2);
     /// ```
     #[inline]
-    pub fn notify_relaxed(&self, n: usize) {
+    pub fn notify_relaxed(&self, n: usize) -> usize {
         self.notify(n.relaxed())
     }
 
@@ -538,7 +540,7 @@ impl Event<()> {
     /// event.notify_additional(1);
     /// ```
     #[inline]
-    pub fn notify_additional(&self, n: usize) {
+    pub fn notify_additional(&self, n: usize) -> usize {
         self.notify(n.additional())
     }
 
@@ -592,7 +594,7 @@ impl Event<()> {
     /// event.notify_additional_relaxed(1);
     /// ```
     #[inline]
-    pub fn notify_additional_relaxed(&self, n: usize) {
+    pub fn notify_additional_relaxed(&self, n: usize) -> usize {
         self.notify(n.additional().relaxed())
     }
 }

diff --git a/src/no_std.rs b/src/no_std.rs
@@ -110,11 +110,11 @@ impl<T> crate::Inner<T> {
 
     /// Notifies a number of entries.
     #[cold]
-    pub(crate) fn notify(&self, mut notify: impl Notification<Tag = T>) {
+    pub(crate) fn notify(&self, mut notify: impl Notification<Tag = T>) -> usize {
         match self.try_lock() {
             Some(mut guard) => {
                 // Notify the listeners.
-                guard.notify(notify);
+                guard.notify(notify)
             }
 
             None => {
@@ -140,6 +140,9 @@ impl<T> crate::Inner<T> {
                 ));
 
                 self.list.queue.push(node);
+
+                // We haven't notified anyone yet.
+                0
             }
         }
     }
@@ -556,23 +559,24 @@ impl<T> ListenerSlab<T> {
 
     /// Notifies a number of listeners.
     #[cold]
-    pub(crate) fn notify(&mut self, mut notify: impl Notification<Tag = T>) {
+    pub(crate) fn notify(&mut self, mut notify: impl Notification<Tag = T>) -> usize {
         let mut n = notify.count(Internal::new());
         let is_additional = notify.is_additional(Internal::new());
         if !is_additional {
             // Make sure we're not notifying more than we have.
             if n <= self.notified {
-                return;
+                return 0;
             }
             n -= self.notified;
         }
 
+        let original_count = n;
         while n > 0 {
             n -= 1;
 
             // Notify the next entry.
             match self.start {
-                None => break,
+                None => return original_count - n - 1,
 
                 Some(e) => {
                     // Get the entry and move the pointer forwards.
@@ -593,6 +597,8 @@ impl<T> ListenerSlab<T> {
                 }
             }
         }
+
+        original_count - n
     }
 
     /// Register a task to be notified when the event is triggered.

diff --git a/src/std.rs b/src/std.rs
@@ -110,7 +110,7 @@ impl<T> crate::Inner<T> {
 
     /// Notifies a number of entries.
     #[cold]
-    pub(crate) fn notify(&self, notify: impl Notification<Tag = T>) {
+    pub(crate) fn notify(&self, notify: impl Notification<Tag = T>) -> usize {
         self.lock().notify(notify)
     }
 
@@ -236,23 +236,24 @@ impl<T> Inner<T> {
     }
 
     #[cold]
-    fn notify(&mut self, mut notify: impl Notification<Tag = T>) {
+    fn notify(&mut self, mut notify: impl Notification<Tag = T>) -> usize {
         let mut n = notify.count(Internal::new());
         let is_additional = notify.is_additional(Internal::new());
 
         if !is_additional {
             if n < self.notified {
-                return;
+                return 0;
             }
             n -= self.notified;
         }
 
+        let original_count = n;
         while n > 0 {
             n -= 1;
 
             // Notify the next entry.
             match self.next {
-                None => break,
+                None => return original_count - n - 1,
 
                 Some(e) => {
                     // Get the entry and move the pointer forwards.
@@ -273,6 +274,8 @@ impl<T> Inner<T> {
                 }
             }
         }
+
+        original_count - n
     }
 }
 

diff --git a/tests/notify.rs b/tests/notify.rs
@@ -24,8 +24,8 @@ fn notify() {
     assert!(!is_notified(l2.as_mut()));
     assert!(!is_notified(l3.as_mut()));
 
-    event.notify(2);
-    event.notify(1);
+    assert_eq!(event.notify(2), 2);
+    assert_eq!(event.notify(1), 0);
 
     assert!(is_notified(l1.as_mut()));
     assert!(is_notified(l2.as_mut()));
@@ -40,9 +40,9 @@ fn notify_additional() {
     let mut l2 = event.listen();
     let mut l3 = event.listen();
 
-    event.notify_additional(1);
-    event.notify(1);
-    event.notify_additional(1);
+    assert_eq!(event.notify_additional(1), 1);
+    assert_eq!(event.notify(1), 0);
+    assert_eq!(event.notify_additional(1), 1);
 
     assert!(is_notified(l1.as_mut()));
     assert!(is_notified(l2.as_mut()));
@@ -59,11 +59,11 @@ fn notify_one() {
     assert!(!is_notified(l1.as_mut()));
     assert!(!is_notified(l2.as_mut()));
 
-    event.notify(1);
+    assert_eq!(event.notify(1), 1);
     assert!(is_notified(l1.as_mut()));
     assert!(!is_notified(l2.as_mut()));
 
-    event.notify(1);
+    assert_eq!(event.notify(1), 1);
     assert!(is_notified(l2.as_mut()));
 }
 
@@ -77,7 +77,7 @@ fn notify_all() {
     assert!(!is_notified(l1.as_mut()));
     assert!(!is_notified(l2.as_mut()));
 
-    event.notify(usize::MAX);
+    assert_eq!(event.notify(usize::MAX), 2);
     assert!(is_notified(l1.as_mut()));
     assert!(is_notified(l2.as_mut()));
 }
@@ -90,7 +90,7 @@ fn drop_notified() {
     let mut l2 = event.listen();
     let mut l3 = event.listen();
 
-    event.notify(1);
+    assert_eq!(event.notify(1), 1);
     drop(l1);
     assert!(is_notified(l2.as_mut()));
     assert!(!is_notified(l3.as_mut()));
@@ -104,7 +104,7 @@ fn drop_notified2() {
     let mut l2 = event.listen();
     let mut l3 = event.listen();
 
-    event.notify(2);
+    assert_eq!(event.notify(2), 2);
     drop(l1);
     assert!(is_notified(l2.as_mut()));
     assert!(!is_notified(l3.as_mut()));
@@ -119,8 +119,8 @@ fn drop_notified_additional() {
     let mut l3 = event.listen();
     let mut l4 = event.listen();
 
-    event.notify_additional(1);
-    event.notify(2);
+    assert_eq!(event.notify_additional(1), 1);
+    assert_eq!(event.notify(2), 1);
     drop(l1);
     assert!(is_notified(l2.as_mut()));
     assert!(is_notified(l3.as_mut()));
@@ -135,7 +135,7 @@ fn drop_non_notified() {
     let mut l2 = event.listen();
     let l3 = event.listen();
 
-    event.notify(1);
+    assert_eq!(event.notify(1), 1);
     drop(l3);
     assert!(is_notified(l1.as_mut()));
     assert!(!is_notified(l2.as_mut()));
@@ -173,7 +173,7 @@ fn notify_all_fair() {
         .poll(&mut Context::from_waker(&waker3))
         .is_pending());
 
-    event.notify(usize::MAX);
+    assert_eq!(event.notify(usize::MAX), 3);
     assert_eq!(&*v.lock().unwrap(), &[1, 2, 3]);
 
     assert!(Pin::new(&mut l1)