From 59180c75b42156c644b9c20b4250f7fa571456e7 Mon Sep 17 00:00:00 2001
From: Paraworker <paraworker@gmail.com>
Date: Wed, 8 Jan 2025 07:21:14 +0000
Subject: [PATCH] Add `WeakLoopHandle`

---
 src/loop_logic.rs | 72 +++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 64 insertions(+), 8 deletions(-)

diff --git a/src/loop_logic.rs b/src/loop_logic.rs
index a1c5ebaa..9584a2f3 100644
--- a/src/loop_logic.rs
+++ b/src/loop_logic.rs
@@ -1,6 +1,6 @@
 use std::cell::{Cell, RefCell};
 use std::fmt::Debug;
-use std::rc::Rc;
+use std::rc::{Rc, Weak};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
@@ -49,15 +49,15 @@ impl RegistrationToken {
 
 pub(crate) struct LoopInner<'l, Data> {
     pub(crate) poll: RefCell<Poll>,
-    // The `Option` is used to keep slots of the slab occipied, to prevent id reuse
-    // while in-flight events might still referr to a recently destroyed event source.
+    // The `Option` is used to keep slots of the slab occupied, to prevent id reuse
+    // while in-flight events might still refer to a recently destroyed event source.
     pub(crate) sources: RefCell<SourceList<'l, Data>>,
     pub(crate) sources_with_additional_lifecycle_events: RefCell<AdditionalLifecycleEventsSet>,
     idles: RefCell<Vec<IdleCallback<'l, Data>>>,
     pending_action: Cell<PostAction>,
 }
 
-/// An handle to an event loop
+/// A handle to an event loop
 ///
 /// This handle allows you to insert new sources and idles in this event loop,
 /// it can be cloned, and it is possible to insert new sources from within a source
@@ -66,7 +66,13 @@ pub struct LoopHandle<'l, Data> {
     inner: Rc<LoopInner<'l, Data>>,
 }
 
-impl<'l, Data> std::fmt::Debug for LoopHandle<'l, Data> {
+/// Weak variant of a [`LoopHandle`]
+#[derive(Clone, Debug, Default)]
+pub struct WeakLoopHandle<'l, Data> {
+    inner: Weak<LoopInner<'l, Data>>,
+}
+
+impl<'l, Data> Debug for LoopHandle<'l, Data> {
     #[cfg_attr(feature = "nightly_coverage", coverage(off))]
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.write_str("LoopHandle { ... }")
@@ -288,6 +294,27 @@ impl<'l, Data> LoopHandle<'l, Data> {
     pub fn adapt_io<F: AsFd>(&self, fd: F) -> crate::Result<crate::io::Async<'l, F>> {
         crate::io::Async::new(self.inner.clone(), fd)
     }
+
+    /// Create a weak reference to this loop data.
+    pub fn downgrade(&self) -> WeakLoopHandle<'l, Data> {
+        WeakLoopHandle {
+            inner: Rc::downgrade(&self.inner),
+        }
+    }
+}
+
+impl<'l, Data> WeakLoopHandle<'l, Data> {
+    /// Try to get a [`LoopHandle`] from this weak reference.
+    ///
+    /// Returns [`None`] if the loop data has been dropped.
+    pub fn upgrade(&self) -> Option<LoopHandle<'l, Data>> {
+        self.inner.upgrade().map(|inner| LoopHandle { inner })
+    }
+
+    /// Check if the loop data has been dropped.
+    pub fn expired(&self) -> bool {
+        self.inner.weak_count() == 0
+    }
 }
 
 /// An event loop
@@ -302,7 +329,7 @@ pub struct EventLoop<'l, Data> {
     synthetic_events: Vec<PollEvent>,
 }
 
-impl<'l, Data> std::fmt::Debug for EventLoop<'l, Data> {
+impl<'l, Data> Debug for EventLoop<'l, Data> {
     #[cfg_attr(feature = "nightly_coverage", coverage(off))]
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.write_str("EventLoop { ... }")
@@ -739,7 +766,7 @@ pub struct LoopSignal {
     notifier: Notifier,
 }
 
-impl std::fmt::Debug for LoopSignal {
+impl Debug for LoopSignal {
     #[cfg_attr(feature = "nightly_coverage", coverage(off))]
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.write_str("LoopSignal { ... }")
@@ -775,7 +802,7 @@ mod tests {
     use crate::{
         channel::{channel, Channel},
         ping::*,
-        EventIterator, EventSource, Poll, PostAction, Readiness, RegistrationToken, Token,
+        timer, EventIterator, EventSource, Poll, PostAction, Readiness, RegistrationToken, Token,
         TokenFactory,
     };
 
@@ -1687,4 +1714,33 @@ mod tests {
             Ok(())
         }
     }
+
+    #[test]
+    fn weak_loop_handle() {
+        let mut event_loop: EventLoop<()> = EventLoop::try_new().unwrap();
+        let weak_handle1 = event_loop.handle().downgrade();
+        let weak_handle2 = weak_handle1.clone();
+        let weak_handle3 = weak_handle1.clone();
+
+        event_loop
+            .handle()
+            .insert_source(timer::Timer::immediate(), move |_, _, _| {
+                // Store and use weak handle in its callback
+                assert!(weak_handle1.upgrade().is_some());
+                timer::TimeoutAction::Drop
+            })
+            .unwrap();
+
+        event_loop.handle().insert_idle(move |_| {
+            // Store and use weak handle in its idle callback
+            assert!(weak_handle2.upgrade().is_some());
+        });
+
+        event_loop.dispatch(None, &mut ()).unwrap();
+
+        drop(event_loop);
+
+        // Test if memory is freed
+        assert!(weak_handle3.expired());
+    }
 }