From b0f6b69b813aae1b7525d222ca1d2ba9c1fa25f1 Mon Sep 17 00:00:00 2001 From: Orson Peters Date: Wed, 28 May 2025 14:39:51 +0200 Subject: [PATCH 1/8] Do not move thread-locals before dropping --- .../std/src/sys/thread_local/native/lazy.rs | 84 ++++++++++++------- 1 file changed, 52 insertions(+), 32 deletions(-) diff --git a/library/std/src/sys/thread_local/native/lazy.rs b/library/std/src/sys/thread_local/native/lazy.rs index 51294285ba013..0cb7fa0ef2481 100644 --- a/library/std/src/sys/thread_local/native/lazy.rs +++ b/library/std/src/sys/thread_local/native/lazy.rs @@ -1,9 +1,9 @@ -use crate::cell::UnsafeCell; -use crate::hint::unreachable_unchecked; +use crate::cell::{Cell, UnsafeCell}; +use crate::mem::MaybeUninit; use crate::ptr; use crate::sys::thread_local::{abort_on_dtor_unwind, destructors}; -pub unsafe trait DestroyedState: Sized { +pub unsafe trait DestroyedState: Sized + Copy { fn register_dtor(s: &Storage); } @@ -19,15 +19,18 @@ unsafe impl DestroyedState for () { } } -enum State { - Initial, - Alive(T), +#[derive(Copy, Clone)] +enum State { + Uninitialized, + Initializing, + Alive, Destroyed(D), } #[allow(missing_debug_implementations)] pub struct Storage { - state: UnsafeCell>, + state: Cell>, + value: UnsafeCell>, } impl Storage @@ -35,7 +38,10 @@ where D: DestroyedState, { pub const fn new() -> Storage { - Storage { state: UnsafeCell::new(State::Initial) } + Storage { + state: Cell::new(State::Uninitialized), + value: UnsafeCell::new(MaybeUninit::uninit()), + } } /// Gets a pointer to the TLS value, potentially initializing it with the @@ -49,35 +55,45 @@ where /// The `self` reference must remain valid until the TLS destructor is run. #[inline] pub unsafe fn get_or_init(&self, i: Option<&mut Option>, f: impl FnOnce() -> T) -> *const T { - let state = unsafe { &*self.state.get() }; - match state { - State::Alive(v) => v, - State::Destroyed(_) => ptr::null(), - State::Initial => unsafe { self.initialize(i, f) }, + if let State::Alive = self.state.get() { + self.value.get().cast() + } else { + self.get_or_init_slow(i, f) } } #[cold] - unsafe fn initialize(&self, i: Option<&mut Option>, f: impl FnOnce() -> T) -> *const T { - // Perform initialization - - let v = i.and_then(Option::take).unwrap_or_else(f); + fn get_or_init_slow(&self, i: Option<&mut Option>, f: impl FnOnce() -> T) -> *const T { + // Ensure we have unique access to an uninitialized value. + match self.state.get() { + State::Uninitialized => self.state.set(State::Initializing), + State::Initializing => panic!("thread_local initializer recursively depends on itself"), + State::Alive => return self.value.get().cast(), + State::Destroyed(_) => return ptr::null(), + } - let old = unsafe { self.state.get().replace(State::Alive(v)) }; - match old { - // If the variable is not being recursively initialized, register - // the destructor. This might be a noop if the value does not need - // destruction. - State::Initial => D::register_dtor(self), - // Else, drop the old value. This might be changed to a panic. - val => drop(val), + struct BackToUninitOnPanic<'a, D>(&'a Cell>); + impl<'a, D> Drop for BackToUninitOnPanic<'a, D> { + fn drop(&mut self) { + self.0.set(State::Uninitialized); + } } - // SAFETY: the state was just set to `Alive` + // Get the initial value, making sure that we restore the state to uninitialized + // should f panic. + let on_panic = BackToUninitOnPanic(&self.state); + let v = i.and_then(Option::take).unwrap_or_else(f); + crate::mem::forget(on_panic); + + // SAFETY: we are !Sync so we have exclusive access to self.value. We also ensured + // that the state was uninitialized so we aren't replacing a value we must keep alive. unsafe { - let State::Alive(v) = &*self.state.get() else { unreachable_unchecked() }; - v + self.value.get().write(MaybeUninit::new(v)); } + + self.state.set(State::Alive); + D::register_dtor(self); + self.value.get().cast() } } @@ -92,9 +108,13 @@ unsafe extern "C" fn destroy(ptr: *mut u8) { // Print a nice abort message if a panic occurs. abort_on_dtor_unwind(|| { let storage = unsafe { &*(ptr as *const Storage) }; - // Update the state before running the destructor as it may attempt to - // access the variable. - let val = unsafe { storage.state.get().replace(State::Destroyed(())) }; - drop(val); + if let State::Alive = storage.state.replace(State::Destroyed(())) { + // SAFETY: we ensured the state was Alive, and prevented running the destructor + // twice by updating the state to Destroyed. This is necessary as the destructor + // may attempt to access the variable. + unsafe { + crate::ptr::drop_in_place(storage.value.get().cast::()); + } + } }) } From f70cf59fc19b7717397e9701b4783f744983275f Mon Sep 17 00:00:00 2001 From: Orson Peters Date: Wed, 28 May 2025 14:51:52 +0200 Subject: [PATCH 2/8] Improve safety comment, double-drop is not relevant here --- library/std/src/sys/thread_local/native/lazy.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/library/std/src/sys/thread_local/native/lazy.rs b/library/std/src/sys/thread_local/native/lazy.rs index 0cb7fa0ef2481..7cf2ba5eed84f 100644 --- a/library/std/src/sys/thread_local/native/lazy.rs +++ b/library/std/src/sys/thread_local/native/lazy.rs @@ -109,9 +109,10 @@ unsafe extern "C" fn destroy(ptr: *mut u8) { abort_on_dtor_unwind(|| { let storage = unsafe { &*(ptr as *const Storage) }; if let State::Alive = storage.state.replace(State::Destroyed(())) { - // SAFETY: we ensured the state was Alive, and prevented running the destructor - // twice by updating the state to Destroyed. This is necessary as the destructor - // may attempt to access the variable. + // SAFETY: we ensured the state was Alive so the value was initialized. + // We also updated the state to Destroyed to prevent the destructor + // from accessing the thread-local variable, as this would violate + // the exclusive access provided by &mut T in Drop::drop. unsafe { crate::ptr::drop_in_place(storage.value.get().cast::()); } From 13bce27e378267203f681470f947208b3e267558 Mon Sep 17 00:00:00 2001 From: Orson Peters Date: Wed, 28 May 2025 16:56:02 +0200 Subject: [PATCH 3/8] Do not panic, maintain old behavior --- .../std/src/sys/thread_local/native/lazy.rs | 42 ++++++++----------- 1 file changed, 17 insertions(+), 25 deletions(-) diff --git a/library/std/src/sys/thread_local/native/lazy.rs b/library/std/src/sys/thread_local/native/lazy.rs index 7cf2ba5eed84f..c5e2618e10f3a 100644 --- a/library/std/src/sys/thread_local/native/lazy.rs +++ b/library/std/src/sys/thread_local/native/lazy.rs @@ -22,7 +22,6 @@ unsafe impl DestroyedState for () { #[derive(Copy, Clone)] enum State { Uninitialized, - Initializing, Alive, Destroyed(D), } @@ -64,36 +63,31 @@ where #[cold] fn get_or_init_slow(&self, i: Option<&mut Option>, f: impl FnOnce() -> T) -> *const T { - // Ensure we have unique access to an uninitialized value. match self.state.get() { - State::Uninitialized => self.state.set(State::Initializing), - State::Initializing => panic!("thread_local initializer recursively depends on itself"), + State::Uninitialized => {} State::Alive => return self.value.get().cast(), State::Destroyed(_) => return ptr::null(), } - struct BackToUninitOnPanic<'a, D>(&'a Cell>); - impl<'a, D> Drop for BackToUninitOnPanic<'a, D> { - fn drop(&mut self) { - self.0.set(State::Uninitialized); - } - } - - // Get the initial value, making sure that we restore the state to uninitialized - // should f panic. - let on_panic = BackToUninitOnPanic(&self.state); let v = i.and_then(Option::take).unwrap_or_else(f); - crate::mem::forget(on_panic); - // SAFETY: we are !Sync so we have exclusive access to self.value. We also ensured - // that the state was uninitialized so we aren't replacing a value we must keep alive. - unsafe { - self.value.get().write(MaybeUninit::new(v)); + match self.state.replace(State::Alive) { + State::Uninitialized => D::register_dtor(self), + + State::Alive => { + // An init occurred during a recursive call, this could be a panic in the future. + + // SAFETY: we cannot be inside a `LocalKey::with` scope, as the initializer + // has already returned and the next scope only starts after we return + // the pointer. Therefore, there can be no references to the old value. + unsafe { (*self.value.get()).assume_init_drop() } + } + + State::Destroyed(_) => unreachable!(), } - self.state.set(State::Alive); - D::register_dtor(self); - self.value.get().cast() + // SAFETY: we are !Sync so we have exclusive access to self.value. + unsafe { (*self.value.get()).write(v) } } } @@ -113,9 +107,7 @@ unsafe extern "C" fn destroy(ptr: *mut u8) { // We also updated the state to Destroyed to prevent the destructor // from accessing the thread-local variable, as this would violate // the exclusive access provided by &mut T in Drop::drop. - unsafe { - crate::ptr::drop_in_place(storage.value.get().cast::()); - } + unsafe { (*storage.value.get()).assume_init_drop() } } }) } From 8785f7b122bbb83d308035565f243ceb95ce4736 Mon Sep 17 00:00:00 2001 From: Orson Peters Date: Wed, 28 May 2025 17:09:15 +0200 Subject: [PATCH 4/8] Add same unsafe bound on get_or_init_slow --- library/std/src/sys/thread_local/native/lazy.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/library/std/src/sys/thread_local/native/lazy.rs b/library/std/src/sys/thread_local/native/lazy.rs index c5e2618e10f3a..2eb1c981edb06 100644 --- a/library/std/src/sys/thread_local/native/lazy.rs +++ b/library/std/src/sys/thread_local/native/lazy.rs @@ -57,12 +57,18 @@ where if let State::Alive = self.state.get() { self.value.get().cast() } else { - self.get_or_init_slow(i, f) + unsafe { self.get_or_init_slow(i, f) } } } + /// # Safety + /// The `self` reference must remain valid until the TLS destructor is run. #[cold] - fn get_or_init_slow(&self, i: Option<&mut Option>, f: impl FnOnce() -> T) -> *const T { + unsafe fn get_or_init_slow( + &self, + i: Option<&mut Option>, + f: impl FnOnce() -> T, + ) -> *const T { match self.state.get() { State::Uninitialized => {} State::Alive => return self.value.get().cast(), From 9ffbc62cb60f2151a38b6c5e18c016e406d10a62 Mon Sep 17 00:00:00 2001 From: Orson Peters Date: Wed, 28 May 2025 17:52:24 +0200 Subject: [PATCH 5/8] When replacing an old value we may not drop it in place --- .../std/src/sys/thread_local/native/lazy.rs | 20 ++++++++----------- 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/library/std/src/sys/thread_local/native/lazy.rs b/library/std/src/sys/thread_local/native/lazy.rs index 2eb1c981edb06..a2bf8d8b968a5 100644 --- a/library/std/src/sys/thread_local/native/lazy.rs +++ b/library/std/src/sys/thread_local/native/lazy.rs @@ -77,23 +77,19 @@ where let v = i.and_then(Option::take).unwrap_or_else(f); + // SAFETY: we cannot be inside a `LocalKey::with` scope, as the initializer + // has already returned and the next scope only starts after we return + // the pointer. Therefore, there can be no references to the old value, + // even if it was initialized. Thus because we are !Sync we have exclusive + // access to self.value and may replace it. + let mut old_value = unsafe { self.value.get().replace(MaybeUninit::new(v)) }; match self.state.replace(State::Alive) { State::Uninitialized => D::register_dtor(self), - - State::Alive => { - // An init occurred during a recursive call, this could be a panic in the future. - - // SAFETY: we cannot be inside a `LocalKey::with` scope, as the initializer - // has already returned and the next scope only starts after we return - // the pointer. Therefore, there can be no references to the old value. - unsafe { (*self.value.get()).assume_init_drop() } - } - + State::Alive => unsafe { old_value.assume_init_drop() }, State::Destroyed(_) => unreachable!(), } - // SAFETY: we are !Sync so we have exclusive access to self.value. - unsafe { (*self.value.get()).write(v) } + self.value.get().cast() } } From 22c5e1d686967d8ef69fec69475f8bbe25b71b2f Mon Sep 17 00:00:00 2001 From: Orson Peters Date: Thu, 29 May 2025 02:47:23 +0200 Subject: [PATCH 6/8] Add test --- .../tls-dont-move-after-init.rs | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 tests/ui/threads-sendsync/tls-dont-move-after-init.rs diff --git a/tests/ui/threads-sendsync/tls-dont-move-after-init.rs b/tests/ui/threads-sendsync/tls-dont-move-after-init.rs new file mode 100644 index 0000000000000..2cd2a88a2c7da --- /dev/null +++ b/tests/ui/threads-sendsync/tls-dont-move-after-init.rs @@ -0,0 +1,39 @@ +//@ run-pass +#![allow(stable_features)] +//@ needs-threads +#![feature(thread_local_try_with)] + +use std::cell::Cell; +use std::thread; + +#[derive(Default)] +struct Foo { + ptr: Cell<*const Foo>, +} + +impl Foo { + fn touch(&self) { + if self.ptr.get().is_null() { + self.ptr.set(self); + } else { + assert!(self.ptr.get() == self); + } + } +} + +impl Drop for Foo { + fn drop(&mut self) { + self.touch(); + } +} + +thread_local!(static FOO: Foo = Foo::default()); + +fn main() { + thread::spawn(|| { + FOO.with(|foo| foo.touch()); + FOO.with(|foo| foo.touch()); + }) + .join() + .unwrap(); +} From aff29df28e312e6ec247a62eaee5c50d431e7015 Mon Sep 17 00:00:00 2001 From: Orson Peters Date: Thu, 29 May 2025 02:48:57 +0200 Subject: [PATCH 7/8] Remove unneeded feature --- tests/ui/threads-sendsync/tls-dont-move-after-init.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/ui/threads-sendsync/tls-dont-move-after-init.rs b/tests/ui/threads-sendsync/tls-dont-move-after-init.rs index 2cd2a88a2c7da..f986202ab89ec 100644 --- a/tests/ui/threads-sendsync/tls-dont-move-after-init.rs +++ b/tests/ui/threads-sendsync/tls-dont-move-after-init.rs @@ -1,7 +1,6 @@ //@ run-pass #![allow(stable_features)] //@ needs-threads -#![feature(thread_local_try_with)] use std::cell::Cell; use std::thread; From b374adc9dbeec2cc969723dbfcb5c8ee990953bd Mon Sep 17 00:00:00 2001 From: Orson Peters Date: Fri, 30 May 2025 12:13:55 +0200 Subject: [PATCH 8/8] Address review comments. --- library/std/src/sys/thread_local/native/lazy.rs | 7 +++++++ tests/ui/threads-sendsync/tls-dont-move-after-init.rs | 1 - 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/library/std/src/sys/thread_local/native/lazy.rs b/library/std/src/sys/thread_local/native/lazy.rs index a2bf8d8b968a5..b556dd9aa25ed 100644 --- a/library/std/src/sys/thread_local/native/lazy.rs +++ b/library/std/src/sys/thread_local/native/lazy.rs @@ -84,8 +84,15 @@ where // access to self.value and may replace it. let mut old_value = unsafe { self.value.get().replace(MaybeUninit::new(v)) }; match self.state.replace(State::Alive) { + // If the variable is not being recursively initialized, register + // the destructor. This might be a noop if the value does not need + // destruction. State::Uninitialized => D::register_dtor(self), + + // Recursive initialization, we only need to drop the old value + // as we've already registered the destructor. State::Alive => unsafe { old_value.assume_init_drop() }, + State::Destroyed(_) => unreachable!(), } diff --git a/tests/ui/threads-sendsync/tls-dont-move-after-init.rs b/tests/ui/threads-sendsync/tls-dont-move-after-init.rs index f986202ab89ec..54fcc32e9bd79 100644 --- a/tests/ui/threads-sendsync/tls-dont-move-after-init.rs +++ b/tests/ui/threads-sendsync/tls-dont-move-after-init.rs @@ -1,5 +1,4 @@ //@ run-pass -#![allow(stable_features)] //@ needs-threads use std::cell::Cell;