Couple of tls handling changes (#120)

* Pass module id to current_tls_addr

* Extract initialize_tls function to share between initialize_main and create

* Extract calculate_tls_size to share between initialize_main and create

The only functional change is that for spawned threads the padding
necessary to correctly align the TLS data is no longer added to the
stack itself and instead becomes unused. Most of the time TLS data will
be aligned no more than the stack alignment, so no memory should be
wasted most of the time.

Author: bjorn3

src/thread/libc.rs

@@ -213,12 +213,12 @@ pub fn errno_location() -> *mut i32 {
     unsafe { __errno_location() }
 }
 
-/// Return the TLS address for the given `offset` for the current thread.
+/// Return the TLS address for the given `module` and `offset` for the current
+/// thread.
 #[inline]
 #[must_use]
-pub fn current_tls_addr(offset: usize) -> *mut c_void {
-    let p = [1, offset];
-    unsafe { __tls_get_addr(&p) }
+pub fn current_tls_addr(module: usize, offset: usize) -> *mut c_void {
+    unsafe { __tls_get_addr(&[module, offset]) }
 }
 
 /// Return the current thread's stack address (lowest address), size, and guard

src/thread/linux_raw.rs

@@ -336,58 +336,96 @@ pub(super) unsafe fn initialize_main(mem: *mut c_void) {
     let canary = random_ptr.read_unaligned();
     __stack_chk_guard = canary;
 
-    let map_size = 0;
+    let mut alloc_size = 0;
+    let (tls_data_bottom, header) = calculate_tls_size(&mut alloc_size);
+
+    // Allocate the thread data. Use `mmap_anonymous` rather than `alloc` here
+    // as the allocator may depend on thread-local data, which is what we're
+    // initializing here.
+    let new = mmap_anonymous(
+        null_mut(),
+        alloc_size,
+        ProtFlags::READ | ProtFlags::WRITE,
+        MapFlags::PRIVATE,
+    )
+    .unwrap()
+    .cast::<u8>();
+
+    let metadata_align = max(unsafe { STARTUP_TLS_INFO.align }, align_of::<Metadata>());
+    debug_assert_eq!(new.addr() % metadata_align, 0);
+
+    let tls_data = new.add(tls_data_bottom);
+    let metadata: *mut Metadata = new.add(header).cast();
+
+    let (newtls, thread_id_ptr) = initialize_tls(
+        tls_data,
+        metadata,
+        canary,
+        stack_least,
+        stack_size,
+        guard_size,
+        0,
+    );
+    let tid = rustix::runtime::set_tid_address(thread_id_ptr.cast());
+    *thread_id_ptr = tid.as_raw_nonzero().get();
+
+    // Point the platform thread-pointer register at the new thread metadata.
+    set_thread_pointer(newtls);
+}
+
+fn calculate_tls_size(map_size: &mut usize) -> (usize, usize) {
+    // SAFETY: `STARTUP_TLS_INFO` is initialized at program startup before
+    // we come here creating new threads.
+    let (startup_tls_align, startup_tls_mem_size) =
+        unsafe { (STARTUP_TLS_INFO.align, STARTUP_TLS_INFO.mem_size) };
 
     // Compute relevant alignments.
-    let tls_data_align = STARTUP_TLS_INFO.align;
+    let tls_data_align = startup_tls_align;
+    let page_align = page_size();
     let header_align = align_of::<Metadata>();
     let metadata_align = max(tls_data_align, header_align);
+    debug_assert!(metadata_align <= page_align);
 
-    // Compute the size to allocate for thread data.
-    let mut alloc_size = 0;
+    *map_size = round_up(*map_size, metadata_align);
 
     // Variant II: TLS data goes below the TCB.
     #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
-    let tls_data_bottom = alloc_size;
+    let tls_data_bottom = *map_size;
 
     #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
     {
-        alloc_size += round_up(STARTUP_TLS_INFO.mem_size, metadata_align);
+        *map_size += round_up(startup_tls_mem_size, tls_data_align);
     }
 
-    let header = alloc_size;
+    let header = *map_size;
 
-    alloc_size += size_of::<Metadata>();
+    *map_size += size_of::<Metadata>();
 
     // Variant I: TLS data goes above the TCB.
     #[cfg(any(target_arch = "aarch64", target_arch = "arm", target_arch = "riscv64"))]
     {
-        alloc_size = round_up(alloc_size, tls_data_align);
+        *map_size = round_up(*map_size, tls_data_align);
     }
 
     #[cfg(any(target_arch = "aarch64", target_arch = "arm", target_arch = "riscv64"))]
-    let tls_data_bottom = alloc_size;
+    let tls_data_bottom = *map_size;
 
     #[cfg(any(target_arch = "aarch64", target_arch = "arm", target_arch = "riscv64"))]
     {
-        alloc_size += round_up(STARTUP_TLS_INFO.mem_size, tls_data_align);
+        *map_size += round_up(startup_tls_mem_size, tls_data_align);
     }
+    (tls_data_bottom, header)
+}
 
-    // Allocate the thread data. Use `mmap_anonymous` rather than `alloc` here
-    // as the allocator may depend on thread-local data, which is what we're
-    // initializing here.
-    let new = mmap_anonymous(
-        null_mut(),
-        alloc_size,
-        ProtFlags::READ | ProtFlags::WRITE,
-        MapFlags::PRIVATE,
-    )
-    .unwrap()
-    .cast::<u8>();
-    debug_assert_eq!(new.addr() % metadata_align, 0);
-
-    let tls_data = new.add(tls_data_bottom);
-    let metadata: *mut Metadata = new.add(header).cast();
+unsafe fn initialize_tls(
+    tls_data: *mut u8,
+    metadata: *mut Metadata,
+    canary: usize,
+    stack_least: *mut u8,
+    stack_size: usize,
+    guard_size: usize,
+    map_size: usize,
+) -> (*mut c_void, *mut i32) {
     let newtls: *mut c_void = (*metadata).abi.thread_pointee.as_mut_ptr().cast();
 
     // Initialize the thread metadata.
@@ -419,11 +457,8 @@ pub(super) unsafe fn initialize_main(mem: *mut c_void) {
     .fill(0);
 
     let thread_id_ptr = (*metadata).thread.thread_id.as_ptr();
-    let tid = rustix::runtime::set_tid_address(thread_id_ptr.cast());
-    *thread_id_ptr = tid.as_raw_nonzero().get();
 
-    // Point the platform thread-pointer register at the new thread metadata.
-    set_thread_pointer(newtls);
+    (newtls, thread_id_ptr)
 }
 
 /// Creates a new thread.
@@ -442,19 +477,9 @@ pub unsafe fn create(
     stack_size: usize,
     guard_size: usize,
 ) -> io::Result<Thread> {
-    // SAFETY: `STARTUP_TLS_INFO` is initialized at program startup before
-    // we come here creating new threads.
-    let (startup_tls_align, startup_tls_mem_size) =
-        unsafe { (STARTUP_TLS_INFO.align, STARTUP_TLS_INFO.mem_size) };
-
     // Compute relevant alignments.
-    let tls_data_align = startup_tls_align;
     let page_align = page_size();
     let stack_align = 16;
-    let header_align = align_of::<Metadata>();
-    let metadata_align = max(tls_data_align, header_align);
-    let stack_metadata_align = max(stack_align, metadata_align);
-    debug_assert!(stack_metadata_align <= page_align);
 
     // Compute the `mmap` size.
     let mut map_size = 0;
@@ -463,36 +488,11 @@ pub unsafe fn create(
 
     let stack_bottom = map_size;
 
-    map_size += round_up(stack_size, stack_metadata_align);
+    map_size += round_up(stack_size, stack_align);
 
     let stack_top = map_size;
 
-    // Variant II: TLS data goes below the TCB.
-    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
-    let tls_data_bottom = map_size;
-
-    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
-    {
-        map_size += round_up(startup_tls_mem_size, tls_data_align);
-    }
-
-    let header = map_size;
-
-    map_size += size_of::<Metadata>();
-
-    // Variant I: TLS data goes above the TCB.
-    #[cfg(any(target_arch = "aarch64", target_arch = "arm", target_arch = "riscv64"))]
-    {
-        map_size = round_up(map_size, tls_data_align);
-    }
-
-    #[cfg(any(target_arch = "aarch64", target_arch = "arm", target_arch = "riscv64"))]
-    let tls_data_bottom = map_size;
-
-    #[cfg(any(target_arch = "aarch64", target_arch = "arm", target_arch = "riscv64"))]
-    {
-        map_size += round_up(startup_tls_mem_size, tls_data_align);
-    }
+    let (tls_data_bottom, header) = calculate_tls_size(&mut map_size);
 
     // Now we'll `mmap` the memory, initialize it, and create the OS thread.
     unsafe {
@@ -519,30 +519,18 @@ pub unsafe fn create(
 
         let tls_data = map.add(tls_data_bottom);
         let metadata: *mut Metadata = map.add(header).cast();
-        let newtls: *mut c_void = (*metadata).abi.thread_pointee.as_mut_ptr().cast();
 
         // Copy the current thread's canary to the new thread.
         let canary = (*current_metadata()).abi.canary;
 
-        // Initialize the thread metadata.
-        metadata.write(Metadata {
-            abi: Abi {
-                canary,
-                dtv: null(),
-                #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
-                this: newtls,
-                _pad: Default::default(),
-                thread_pointee: [],
-            },
-            thread: ThreadData::new(stack_least.cast(), stack_size, guard_size, map_size),
-        });
-
-        // Initialize the TLS data with explicit initializer data.
-        slice::from_raw_parts_mut(tls_data, STARTUP_TLS_INFO.file_size).copy_from_slice(
-            slice::from_raw_parts(
-                STARTUP_TLS_INFO.addr.cast::<u8>(),
-                STARTUP_TLS_INFO.file_size,
-            ),
+        let (newtls, thread_id_ptr) = initialize_tls(
+            tls_data,
+            metadata,
+            canary,
+            stack_least,
+            stack_size,
+            guard_size,
+            map_size,
         );
 
         // Allocate space for the thread arguments on the child's stack.
@@ -582,7 +570,6 @@ pub unsafe fn create(
             | CloneFlags::CHILD_CLEARTID
             | CloneFlags::CHILD_SETTID
             | CloneFlags::PARENT_SETTID;
-        let thread_id_ptr = (*metadata).thread.thread_id.as_ptr();
         let clone_res = clone(
             flags.bits(),
             stack.cast(),
@@ -1005,10 +992,15 @@ pub fn errno_location() -> *mut i32 {
     unsafe { core::ptr::addr_of_mut!((*current_metadata()).thread.errno_val).cast::<i32>() }
 }
 
-/// Return the TLS address for the given `offset` for the current thread.
+/// Return the TLS address for the given `module` and `offset` for the current
+/// thread.
 #[inline]
 #[must_use]
-pub fn current_tls_addr(offset: usize) -> *mut c_void {
+pub fn current_tls_addr(module: usize, offset: usize) -> *mut c_void {
+    // Offset 0 is the generation field, and we don't support dynamic linking,
+    // so we should only ever see 1 here.
+    assert_eq!(module, 1);
+
     // Platforms where TLS data goes after the ABI-exposed fields.
     #[cfg(any(target_arch = "aarch64", target_arch = "arm", target_arch = "riscv64"))]
     {