From 9ecb8e614f107f68b5c6ba770342ae72af1cd07b Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Sun, 1 Aug 2021 09:16:11 -0700
Subject: [PATCH] feat: improve panics on thread ID bit exhaustion

Signed-off-by: Eliza Weisman <eliza@buoyant.io>
---
 src/macros.rs | 19 +++++++++++++++++++
 src/tid.rs    | 29 +++++++++++++++++++++--------
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/src/macros.rs b/src/macros.rs
index d9ccd22..d5b1f58 100644
--- a/src/macros.rs
+++ b/src/macros.rs
@@ -23,3 +23,22 @@ macro_rules! test_dbg {
         }
     };
 }
+
+macro_rules! panic_in_drop {
+    ($($arg:tt)*) => {
+        if !std::thread::panicking() {
+            panic!($($arg)*)
+        } else {
+            let thread = std::thread::current();
+            eprintln!(
+                "thread '{thread}' attempted to panic at '{msg}', {file}:{line}:{col}\n\
+                note: we were already unwinding due to a previous panic.",
+                thread = thread.name().unwrap_or("<unnamed>"),
+                msg = format_args!($($arg)*),
+                file = file!(),
+                line = line!(),
+                col = column!(),
+            );
+        }
+    }
+}
diff --git a/src/tid.rs b/src/tid.rs
index f410e6f..57d64f9 100644
--- a/src/tid.rs
+++ b/src/tid.rs
@@ -12,6 +12,7 @@ use std::{
     collections::VecDeque,
     fmt,
     marker::PhantomData,
+    sync::PoisonError,
 };
 
 /// Uniquely identifies a thread.
@@ -138,10 +139,10 @@ impl Registration {
     #[inline(always)]
     fn current<C: cfg::Config>(&self) -> Tid<C> {
         if let Some(tid) = self.0.get().map(Tid::new) {
-            tid
-        } else {
-            self.register()
+            return tid;
         }
+
+        self.register()
     }
 
     #[cold]
@@ -157,8 +158,21 @@ impl Registration {
                     None
                 }
             })
-            .unwrap_or_else(|| REGISTRY.next.fetch_add(1, Ordering::AcqRel));
-        debug_assert!(id <= Tid::<C>::BITS, "thread ID overflow!");
+            .unwrap_or_else(|| {
+                let id = REGISTRY.next.fetch_add(1, Ordering::AcqRel);
+                if id > Tid::<C>::BITS {
+                    panic_in_drop!(
+                        "creating a new thread ID ({}) would exceed the \
+                        maximum number of thread ID bits specified in {} \
+                        ({})",
+                        id,
+                        std::any::type_name::<C>(),
+                        Tid::<C>::BITS,
+                    );
+                }
+                id
+            });
+
         self.0.set(Some(id));
         Tid::new(id)
     }
@@ -173,9 +187,8 @@ impl Registration {
 impl Drop for Registration {
     fn drop(&mut self) {
         if let Some(id) = self.0.get() {
-            if let Ok(mut free) = REGISTRY.free.lock() {
-                free.push_back(id);
-            }
+            let mut free_list = REGISTRY.free.lock().unwrap_or_else(PoisonError::into_inner);
+            free_list.push_back(id);
         }
     }
 }