From 0d0fa4f2dc0ab03fa3f58f465c6147e606a3b626 Mon Sep 17 00:00:00 2001
From: Mingwei Samuel <mingwei.samuel@gmail.com>
Date: Tue, 26 Mar 2024 21:09:37 -0700
Subject: [PATCH] fix: potential race condition in `Notified` `Drop` impl

---
 riven/src/util/notify.rs | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/riven/src/util/notify.rs b/riven/src/util/notify.rs
index 1a34005..74734ec 100644
--- a/riven/src/util/notify.rs
+++ b/riven/src/util/notify.rs
@@ -6,6 +6,8 @@ use futures::future::FusedFuture;
 use futures::FutureExt;
 use parking_lot::Mutex;
 use slab::Slab;
+#[cfg(feature = "tracing")]
+use tracing as log;
 
 #[derive(Default)]
 pub struct Notify {
@@ -64,10 +66,10 @@ impl Notify {
                     // Ensure generation matches before removing, to prevent ABA problem.
                     // If no match, means `Notify::notify_waiters` has already been called and deallocated us.
                     if internal.generation == generation {
-                        internal
-                            .waiters
-                            .try_remove(key)
-                            .expect("Expected to drop registered `Notified` but waker not found.");
+                        if internal.waiters.try_remove(key).is_none() {
+                            // Rare race condition to get here, maybe?
+                            log::trace!("Tried to de-register `Notified` on drop but corresponding waker not found.");
+                        }
                         internal.waiters.shrink_to_fit();
                     }
                 }