From a7e6ace4a08f2d2b405a9ab88a5766ef593e16e5 Mon Sep 17 00:00:00 2001
From: Armin Ronacher <armin.ronacher@active-4.com>
Date: Tue, 18 Oct 2022 19:41:56 +0200
Subject: [PATCH] Sticky drops from the registry eagerly now

---
 CHANGELOG.md    |  2 ++
 src/registry.rs |  9 +++++++++
 src/sticky.rs   | 12 ++++++++++++
 3 files changed, 23 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2319617..e200b76 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,8 @@ All notable changes to similar are documented here.
 * `Fragile` no longer boxes internally.
 * `Sticky` and `SemiSticky` now require the use of stack tokens.
   For more information see [#26](https://github.com/mitsuhiko/fragile/issues/26)
+* `Sticky` now tries to drop entries from the thread local registry eagerly
+  if it's dropped on the right thread.
 
 ## 1.2.1
 
diff --git a/src/registry.rs b/src/registry.rs
index 663310f..1ee070d 100644
--- a/src/registry.rs
+++ b/src/registry.rs
@@ -34,6 +34,11 @@ mod slab_impl {
         let _ = thread_id;
         REGISTRY.with(|registry| unsafe { (*registry.get()).0.remove(item_id) })
     }
+
+    pub fn try_remove(item_id: ItemId, thread_id: NonZeroUsize) -> Option<Entry> {
+        let _ = thread_id;
+        REGISTRY.with(|registry| unsafe { (*registry.get()).0.try_remove(item_id) })
+    }
 }
 
 #[cfg(not(feature = "slab"))]
@@ -76,6 +81,10 @@ mod map_impl {
         REGISTRY
             .with(|registry| unsafe { (*registry.get()).0.remove(&(thread_id, item_id)).unwrap() })
     }
+
+    pub fn try_remove(item_id: ItemId, thread_id: NonZeroUsize) -> Option<Entry> {
+        REGISTRY.with(|registry| unsafe { (*registry.get()).0.remove(&(thread_id, item_id)) })
+    }
 }
 
 #[cfg(feature = "slab")]
diff --git a/src/sticky.rs b/src/sticky.rs
index 291ab8e..bc15c40 100644
--- a/src/sticky.rs
+++ b/src/sticky.rs
@@ -35,12 +35,24 @@ pub struct Sticky<T: 'static> {
 
 impl<T> Drop for Sticky<T> {
     fn drop(&mut self) {
+        // if the type needs dropping we can only do so on the
+        // right thread.  worst case we leak the value until the
+        // thread dies.
         if mem::needs_drop::<T>() {
             unsafe {
                 if self.is_valid() {
                     self.unsafe_take_value();
                 }
             }
+
+        // otherwise we take the liberty to drop the value
+        // right here and now.  We can however only do that if
+        // we are on the right thread.  If we are not, we again
+        // need to wait for the thread to shut down.
+        } else if let Some(entry) = registry::try_remove(self.item_id, self.thread_id) {
+            unsafe {
+                (entry.drop)(entry.ptr);
+            }
         }
     }
 }