159
159
mod tests;
160
160
161
161
use crate :: any:: Any ;
162
- use crate :: cell:: { OnceCell , UnsafeCell } ;
162
+ use crate :: cell:: { Cell , OnceCell , UnsafeCell } ;
163
163
use crate :: env;
164
164
use crate :: ffi:: { CStr , CString } ;
165
165
use crate :: fmt;
@@ -699,16 +699,18 @@ where
699
699
700
700
thread_local ! {
701
701
static CURRENT : OnceCell <Thread > = const { OnceCell :: new( ) } ;
702
+ static CURRENT_ID : Cell <Option <ThreadId >> = const { Cell :: new( None ) } ;
702
703
}
703
704
704
705
/// Sets the thread handle for the current thread.
705
706
///
706
707
/// Aborts if the handle has been set already to reduce code size.
707
708
pub ( crate ) fn set_current ( thread : Thread ) {
709
+ let tid = thread. id ( ) ;
708
710
// Using `unwrap` here can add ~3kB to the binary size. We have complete
709
711
// control over where this is called, so just abort if there is a bug.
710
712
CURRENT . with ( |current| match current. set ( thread) {
711
- Ok ( ( ) ) => { }
713
+ Ok ( ( ) ) => CURRENT_ID . set ( Some ( tid ) ) ,
712
714
Err ( _) => rtabort ! ( "thread::set_current should only be called once per thread" ) ,
713
715
} ) ;
714
716
}
@@ -718,7 +720,25 @@ pub(crate) fn set_current(thread: Thread) {
718
720
/// In contrast to the public `current` function, this will not panic if called
719
721
/// from inside a TLS destructor.
720
722
pub ( crate ) fn try_current ( ) -> Option < Thread > {
721
- CURRENT . try_with ( |current| current. get_or_init ( || Thread :: new_unnamed ( ) ) . clone ( ) ) . ok ( )
723
+ CURRENT
724
+ . try_with ( |current| {
725
+ let thread = current. get_or_init ( Thread :: new_unnamed) . clone ( ) ;
726
+ CURRENT_ID . set ( Some ( thread. id ( ) ) ) ;
727
+ thread
728
+ } )
729
+ . ok ( )
730
+ }
731
+
732
+ /// Gets the id of the thread that invokes it.
733
+ ///
734
+ /// If called from inside a TLS destructor and the thread was never
735
+ /// assigned an id, returns `None`.
736
+ #[ inline]
737
+ pub ( crate ) fn try_current_id ( ) -> Option < ThreadId > {
738
+ if CURRENT_ID . get ( ) . is_none ( ) {
739
+ let _ = try_current ( ) ;
740
+ }
741
+ CURRENT_ID . get ( )
722
742
}
723
743
724
744
/// Gets a handle to the thread that invokes it.
0 commit comments