Fix macos tests

wenyuzhao · Jul 1, 2024 · 88072f2 · 88072f2
1 parent b36208b
commit 88072f2
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 4 deletions.
diff --git a/mallockit/macros/src/lib.rs b/mallockit/macros/src/lib.rs
@@ -48,6 +48,7 @@ pub fn plan(_attr: TokenStream, item: TokenStream) -> TokenStream {
             #[cfg(any(feature = "malloc", feature = "mallockit/malloc"))]
             #[::mallockit::ctor]
             unsafe fn ctor() {
+                <<Plan as ::mallockit::Plan>::Mutator as ::mallockit::mutator::TLS>::current();
                 ::mallockit::util::sys::hooks::process_start(&*PLAN);
             }
 
@@ -92,20 +93,30 @@ pub fn mutator(_attr: TokenStream, item: TokenStream) -> TokenStream {
 
 
         mod __mallockit_mutator {
+            #[cfg(not(target_os = "macos"))]
             fn init() -> super::#name {
                 ::mallockit::mutator::init_pthread_specific();
                 <super::#name as ::mallockit::Mutator>::new()
             }
 
+            #[cfg(not(target_os = "macos"))]
             #[thread_local]
             pub(super) static mut MUTATOR: ::mallockit::util::Lazy<super::#name, ::mallockit::util::Local> = ::mallockit::util::Lazy::new(init);
 
             #[no_mangle]
+            #[cfg(not(target_os = "macos"))]
             extern "C" fn mallockit_pthread_destructor() {
                 unsafe {
                     MUTATOR.reset(init);
                 }
             }
+
+            #[no_mangle]
+            #[cfg(target_os = "macos")]
+            extern "C" fn mallockit_pthread_destructor() {
+                use crate::mallockit::mutator::TLS;
+                <super::#name as ::mallockit::mutator::TLS>::current().reset();
+            }
         }
 
         impl ::mallockit::mutator::TLS for #name {

diff --git a/mallockit/src/mutator.rs b/mallockit/src/mutator.rs
@@ -100,6 +100,10 @@ pub trait TLS: Sized {
     fn current() -> &'static mut Self {
         unsafe { &mut *macos_tls::get_tls::<Self>() }
     }
+
+    fn reset(&mut self) {
+        *self = Self::new();
+    }
 }
 
 impl TLS for u8 {
@@ -211,8 +215,7 @@ mod macos_tls {
     fn init_tls<T: TLS>() -> *mut (InternalTLS, T) {
         let ptr = alloc_tls::<(InternalTLS, T)>();
         unsafe {
-            (*ptr).0 = InternalTLS::new();
-            (*ptr).1 = T::new();
+            std::ptr::write(&mut (*ptr).0, InternalTLS::new());
             unsafe {
                 let mut tcb: *mut *mut T;
                 asm! {
@@ -224,6 +227,8 @@ mod macos_tls {
                 }
                 tcb.add(SLOT).write(ptr as *mut T)
             }
+            std::ptr::write(&mut (*ptr).1, T::new());
+            crate::mutator::init_pthread_specific();
         }
         ptr
     }
@@ -234,9 +239,10 @@ mod macos_tls {
     fn init_tls<T: TLS>() -> *mut (InternalTLS, T) {
         let ptr = alloc_tls::<(InternalTLS, T)>();
         unsafe {
-            (*ptr).0 = InternalTLS::new();
-            (*ptr).1 = T::new();
+            std::ptr::write(&mut (*ptr).0, InternalTLS::new());
             asm!("mov gs:{offset}, {0}", in(reg) ptr, offset = const OFFSET);
+            std::ptr::write(&mut (*ptr).1, T::new());
+            crate::mutator::init_pthread_specific();
         }
         ptr
     }