Expand test with allocating from thread-local destructors

tests/ui/threads-sendsync/tls-in-global-alloc.rs

@@ -2,27 +2,46 @@
 //@ needs-threads
 
 use std::alloc::{GlobalAlloc, Layout, System};
-use std::thread::Thread;
-use std::sync::atomic::{AtomicUsize, AtomicBool, Ordering};
+use std::hint::black_box;
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::thread::{Thread, ThreadId};
 
 static GLOBAL: AtomicUsize = AtomicUsize::new(0);
+static SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS: AtomicBool = AtomicBool::new(false);
+static LOCAL_TRY_WITH_SUCCEEDED_ALLOC: AtomicBool = AtomicBool::new(false);
+static LOCAL_TRY_WITH_SUCCEEDED_DEALLOC: AtomicBool = AtomicBool::new(false);
 
-struct Local(Thread);
+struct LocalForAllocatorWithoutDrop(ThreadId);
 
-thread_local! {
-    static LOCAL: Local = {
-        GLOBAL.fetch_or(1, Ordering::Relaxed);
-        Local(std::thread::current())
-    };
-}
+struct LocalForAllocatorWithDrop(Thread);
 
-impl Drop for Local {
+impl Drop for LocalForAllocatorWithDrop {
     fn drop(&mut self) {
         GLOBAL.fetch_or(2, Ordering::Relaxed);
     }
 }
 
-static SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS: AtomicBool = AtomicBool::new(false);
+struct LocalForUser(u32);
+
+impl Drop for LocalForUser {
+    // A user might call the global allocator in a thread-local drop.
+    fn drop(&mut self) {
+        self.0 += 1;
+        drop(black_box(Box::new(self.0)))
+    }
+}
+
+thread_local! {
+    static LOCAL_FOR_USER0: LocalForUser = LocalForUser(0);
+    static LOCAL_FOR_ALLOCATOR_WITHOUT_DROP: LocalForAllocatorWithoutDrop = {
+        LocalForAllocatorWithoutDrop(std::thread::current().id())
+    };
+    static LOCAL_FOR_ALLOCATOR_WITH_DROP: LocalForAllocatorWithDrop = {
+        GLOBAL.fetch_or(1, Ordering::Relaxed);
+        LocalForAllocatorWithDrop(std::thread::current())
+    };
+    static LOCAL_FOR_USER1: LocalForUser = LocalForUser(1);
+}
 
 #[global_allocator]
 static ALLOC: Alloc = Alloc;
@@ -33,9 +52,19 @@ unsafe impl GlobalAlloc for Alloc {
         // Make sure we aren't re-entrant.
         assert!(!SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS.load(Ordering::Relaxed));
         SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS.store(true, Ordering::Relaxed);
-        LOCAL.with(|local| {
+
+        // Should be infallible.
+        LOCAL_FOR_ALLOCATOR_WITHOUT_DROP.with(|local| {
+            assert!(local.0 == std::thread::current().id());
+        });
+
+        // May fail once thread-local destructors start running, and ours has
+        // been ran.
+        let try_with_ret = LOCAL_FOR_ALLOCATOR_WITH_DROP.try_with(|local| {
             assert!(local.0.id() == std::thread::current().id());
         });
+        LOCAL_TRY_WITH_SUCCEEDED_ALLOC.fetch_or(try_with_ret.is_ok(), Ordering::Relaxed);
+
         let ret = unsafe { System.alloc(layout) };
         SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS.store(false, Ordering::Relaxed);
         ret
@@ -45,9 +74,19 @@ unsafe impl GlobalAlloc for Alloc {
         // Make sure we aren't re-entrant.
         assert!(!SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS.load(Ordering::Relaxed));
         SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS.store(true, Ordering::Relaxed);
-        LOCAL.with(|local| {
+
+        // Should be infallible.
+        LOCAL_FOR_ALLOCATOR_WITHOUT_DROP.with(|local| {
+            assert!(local.0 == std::thread::current().id());
+        });
+
+        // May fail once thread-local destructors start running, and ours has
+        // been ran.
+        let try_with_ret = LOCAL_FOR_ALLOCATOR_WITH_DROP.try_with(|local| {
             assert!(local.0.id() == std::thread::current().id());
         });
+        LOCAL_TRY_WITH_SUCCEEDED_DEALLOC.fetch_or(try_with_ret.is_ok(), Ordering::Relaxed);
+
         unsafe { System.dealloc(ptr, layout) }
         SHOULD_PANIC_ON_GLOBAL_ALLOC_ACCESS.store(false, Ordering::Relaxed);
     }
@@ -55,10 +94,14 @@ unsafe impl GlobalAlloc for Alloc {
 
 fn main() {
     std::thread::spawn(|| {
+        LOCAL_FOR_USER0.with(|l| assert!(l.0 == 0));
         std::hint::black_box(vec![1, 2]);
         assert!(GLOBAL.load(Ordering::Relaxed) == 1);
+        LOCAL_FOR_USER1.with(|l| assert!(l.0 == 1));
     })
     .join()
     .unwrap();
     assert!(GLOBAL.load(Ordering::Relaxed) == 3);
+    assert!(LOCAL_TRY_WITH_SUCCEEDED_ALLOC.load(Ordering::Relaxed));
+    assert!(LOCAL_TRY_WITH_SUCCEEDED_DEALLOC.load(Ordering::Relaxed));
 }