Upgrade to use snmalloc for rust compiler

Cargo.lock

@@ -658,6 +658,15 @@ dependencies = [
  "rustc-semver",
 ]
 
+[[package]]
+name = "cmake"
+version = "0.1.50"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a31c789563b815f77f4250caee12365734369f942439b7defd71e18a48197130"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "collect-license-metadata"
 version = "0.1.0"
@@ -1147,19 +1156,6 @@ version = "1.0.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "86e3bdc80eee6e16b2b6b0f87fbc98c04bee3455e35174c0de1a125d0688c632"
 
-[[package]]
-name = "dlmalloc"
-version = "0.2.6"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3264b043b8e977326c1ee9e723da2c1f8d09a99df52cacf00b4dbce5ac54414d"
-dependencies = [
- "cfg-if",
- "compiler_builtins",
- "libc",
- "rustc-std-workspace-core",
- "windows-sys 0.52.0",
-]
-
 [[package]]
 name = "either"
 version = "1.10.0"
@@ -1178,6 +1174,12 @@ dependencies = [
  "serde_json",
 ]
 
+[[package]]
+name = "elf"
+version = "0.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4445909572dbd556c457c849c4ca58623d84b27c8fff1e74b0b4227d8b90d17b"
+
 [[package]]
 name = "elsa"
 version = "1.7.1"
@@ -5141,6 +5143,18 @@ version = "1.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b"
 
+[[package]]
+name = "snmalloc-edp"
+version = "0.1.0"
+source = "git+https://github.com/fortanix/rust-sgx?branch=aj/update-sgx-alloc#ac0c2a16324f2a74b5c2d9b6dbf204fb3c89b0f7"
+dependencies = [
+ "cc",
+ "cmake",
+ "compiler_builtins",
+ "elf",
+ "rustc-std-workspace-core",
+]
+
 [[package]]
 name = "socket2"
 version = "0.5.6"
@@ -5232,7 +5246,6 @@ dependencies = [
  "cfg-if",
  "compiler_builtins",
  "core",
- "dlmalloc",
  "fortanix-sgx-abi",
  "hashbrown",
  "hermit-abi",
@@ -5247,6 +5260,7 @@ dependencies = [
  "rand",
  "rand_xorshift",
  "rustc-demangle",
+ "snmalloc-edp",
  "std_detect",
  "unwind",
  "wasi",

library/std/Cargo.toml

@@ -42,7 +42,7 @@ rand = { version = "0.8.5", default-features = false, features = ["alloc"] }
 rand_xorshift = "0.3.0"
 
 [target.'cfg(any(all(target_family = "wasm", target_os = "unknown"), target_os = "xous", all(target_vendor = "fortanix", target_env = "sgx")))'.dependencies]
-dlmalloc = { version = "0.2.4", features = ['rustc-dep-of-std'] }
+snmalloc-edp = { git = "https://github.com/fortanix/rust-sgx", branch = "aj/update-sgx-alloc", features = ['rustc-dep-of-std'], public = true }
 
 [target.x86_64-fortanix-unknown-sgx.dependencies]
 fortanix-sgx-abi = { version = "0.5.0", features = ['rustc-dep-of-std'], public = true }

library/std/src/sys/pal/sgx/abi/entry.S

@@ -97,8 +97,8 @@ IMAGE_BASE:
 .equ tcsls_user_r13,            0x48
 .equ tcsls_user_r14,            0x50
 .equ tcsls_user_r15,            0x58
-.equ tcsls_tls_ptr,             0x60
-.equ tcsls_tcs_addr,            0x68
+.equ tcsls_tcs_addr,            0x60
+.equ tcsls_tls_ptr,             0x68
 
 .macro load_tcsls_flag_secondary_bool reg:req comments:vararg
     .ifne tcsls_flag_secondary /* to convert to a bool, must be the first bit */
@@ -131,7 +131,7 @@ elf_entry:
     sub %rax,%rdx                     /* all chars written?     */
     jnz .Lelf_entry_call
 
-.Lelf_exit:    
+.Lelf_exit:
     movq $60,%rax                     /* exit() syscall         */
     movq $1,%rdi                      /* exit code 1            */
     syscall
@@ -355,14 +355,14 @@ get_tcs_addr:
 
 .global get_tls_ptr
 get_tls_ptr:
-    mov %gs:tcsls_tls_ptr,%rax
+    mov %gs:tcsls_tls_ptr(,%rdi,8),%rax
     pop %r11
     lfence
     jmp *%r11
 
 .global set_tls_ptr
 set_tls_ptr:
-    mov %rdi,%gs:tcsls_tls_ptr
+    mov %rsi,%gs:tcsls_tls_ptr(,%rdi,8)
     pop %r11
     lfence
     jmp *%r11

library/std/src/sys/pal/sgx/abi/mem.rs

@@ -19,11 +19,13 @@ extern "C" {
 }
 
 /// Returns the base memory address of the heap
+#[allow(dead_code)]
 pub(crate) fn heap_base() -> *const u8 {
     unsafe { rel_ptr_mut(HEAP_BASE) }
 }
 
 /// Returns the size of the heap
+#[allow(dead_code)]
 pub(crate) fn heap_size() -> usize {
     unsafe { HEAP_SIZE }
 }

library/std/src/sys/pal/sgx/abi/mod.rs

@@ -2,7 +2,8 @@
 
 use crate::io::Write;
 use core::arch::global_asm;
-use core::sync::atomic::{AtomicUsize, Ordering};
+use core::sync::atomic::{AtomicUsize, Ordering, AtomicBool};
+use snmalloc_edp::*;
 
 // runtime features
 pub(super) mod panic;
@@ -18,12 +19,16 @@ pub mod usercalls;
 #[cfg(not(test))]
 global_asm!(include_str!("entry.S"), options(att_syntax));
 
+// To initialize global allocator once per program
+static INIT: AtomicBool = AtomicBool::new(false);
+
 #[repr(C)]
 struct EntryReturn(u64, u64);
 
 #[cfg(not(test))]
 #[no_mangle]
 unsafe extern "C" fn tcs_init(secondary: bool) {
+
     // Be very careful when changing this code: it runs before the binary has been
     // relocated. Any indirect accesses to symbols will likely fail.
     const UNINIT: usize = 0;
@@ -41,6 +46,12 @@ unsafe extern "C" fn tcs_init(secondary: bool) {
         // This thread just obtained the lock and other threads will observe BUSY
         Ok(_) => {
             reloc::relocate_elf_rela();
+
+            // Snmalloc global allocator initialization
+            if !INIT.swap(true, Ordering::Relaxed) {
+                unsafe { sn_global_init(); }
+            }
+
             RELOC_STATE.store(DONE, Ordering::Release);
         }
         // We need to wait until the initialization is done.
@@ -61,6 +72,14 @@ unsafe extern "C" fn tcs_init(secondary: bool) {
 #[cfg(not(test))]
 #[no_mangle]
 extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64) -> EntryReturn {
+
+    // Initialize thread local allocator
+    let mut allocator = core::mem::MaybeUninit::<snmalloc_edp::Alloc>::uninit();
+    unsafe {
+        sn_thread_init(allocator.as_mut_ptr());
+        tls::set_tls_ptr(tls::TlsIndex::AllocPtr, allocator.as_mut_ptr() as *const u8);
+    }
+
     // FIXME: how to support TLS in library mode?
     let tls = Box::new(tls::Tls::new());
     let tls_guard = unsafe { tls.activate() };
@@ -69,6 +88,9 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
         let join_notifier = super::thread::Thread::entry();
         drop(tls_guard);
         drop(join_notifier);
+        drop(tls);
+
+        alloc_thread_cleanup(allocator.as_mut_ptr());
 
         EntryReturn(0, 0)
     } else {
@@ -87,11 +109,23 @@ extern "C" fn entry(p1: u64, p2: u64, p3: u64, secondary: bool, p4: u64, p5: u64
             // main function, so we pass these in as the standard pointer-sized
             // values in `argc` and `argv`.
             let ret = main(p2 as _, p1 as _);
+            if ret == 0 {
+                drop(tls_guard);
+                drop(tls);
+                alloc_thread_cleanup(allocator.as_mut_ptr());
+            }
             exit_with_code(ret)
         }
     }
 }
 
+pub(super) fn alloc_thread_cleanup(allocator: *mut snmalloc_edp::Alloc) {
+    unsafe {
+        tls::set_tls_ptr(tls::TlsIndex::AllocPtr, core::ptr::null_mut());
+        sn_thread_cleanup(allocator);
+    }
+}
+
 pub(super) fn exit_with_code(code: isize) -> ! {
     if code != 0 {
         if let Some(mut out) = panic::SgxPanicOutput::new() {

library/std/src/sys/pal/sgx/abi/tls/mod.rs

@@ -23,9 +23,17 @@ macro_rules! dup {
 #[export_name = "_ZN16__rust_internals3std3sys3sgx3abi3tls14TLS_DESTRUCTORE"]
 static TLS_DESTRUCTOR: [AtomicUsize; TLS_KEYS] = dup!((* * * * * * *) (AtomicUsize::new(0)));
 
+// Use TCSLS to store both the thread specific allocator
+// and TLS address at different indices
+#[repr(u8)]
+pub enum TlsIndex {
+    AllocPtr = 0,
+    TlsPtr = 1,
+}
+
 extern "C" {
-    fn get_tls_ptr() -> *const u8;
-    fn set_tls_ptr(tls: *const u8);
+    pub fn get_tls_ptr(index: TlsIndex) -> *const u8;
+    pub fn set_tls_ptr(index: TlsIndex, tls: *const u8);
 }
 
 #[derive(Copy, Clone)]
@@ -88,20 +96,20 @@ impl Tls {
 
     pub unsafe fn activate(&self) -> ActiveTls<'_> {
         // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        unsafe { set_tls_ptr(self as *const Tls as _) };
+        unsafe { set_tls_ptr(TlsIndex::TlsPtr, self as *const Tls as _) };
         ActiveTls { tls: self }
     }
 
     #[allow(unused)]
     pub unsafe fn activate_persistent(self: Box<Self>) {
         // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        unsafe { set_tls_ptr(core::ptr::addr_of!(*self) as _) };
+        unsafe { set_tls_ptr(TlsIndex::TlsPtr, core::ptr::addr_of!(*self) as _) };
         mem::forget(self);
     }
 
     unsafe fn current<'a>() -> &'a Tls {
         // FIXME: Needs safety information. See entry.S for `set_tls_ptr` definition.
-        unsafe { &*(get_tls_ptr() as *const Tls) }
+        unsafe { &*(get_tls_ptr(TlsIndex::TlsPtr) as *const Tls) }
     }
 
     pub fn create(dtor: Option<unsafe extern "C" fn(*mut u8)>) -> Key {

library/std/src/sys/pal/sgx/alloc.rs

@@ -1,85 +1,29 @@
-use crate::alloc::{GlobalAlloc, Layout, System};
-use crate::ptr;
-use core::sync::atomic::{AtomicBool, Ordering};
-
-use super::abi::mem as sgx_mem;
-use super::waitqueue::SpinMutex;
-
-// Using a SpinMutex because we never want to exit the enclave waiting for the
-// allocator.
-//
-// The current allocator here is the `dlmalloc` crate which we've got included
-// in the rust-lang/rust repository as a submodule. The crate is a port of
-// dlmalloc.c from C to Rust.
-#[cfg_attr(test, linkage = "available_externally")]
-#[export_name = "_ZN16__rust_internals3std3sys3sgx5alloc8DLMALLOCE"]
-static DLMALLOC: SpinMutex<dlmalloc::Dlmalloc<Sgx>> =
-    SpinMutex::new(dlmalloc::Dlmalloc::new_with_allocator(Sgx {}));
-
-struct Sgx;
-
-unsafe impl dlmalloc::Allocator for Sgx {
-    /// Allocs system resources
-    fn alloc(&self, _size: usize) -> (*mut u8, usize, u32) {
-        static INIT: AtomicBool = AtomicBool::new(false);
-
-        // No ordering requirement since this function is protected by the global lock.
-        if !INIT.swap(true, Ordering::Relaxed) {
-            (sgx_mem::heap_base() as _, sgx_mem::heap_size(), 0)
-        } else {
-            (ptr::null_mut(), 0, 0)
-        }
-    }
-
-    fn remap(&self, _ptr: *mut u8, _oldsize: usize, _newsize: usize, _can_move: bool) -> *mut u8 {
-        ptr::null_mut()
-    }
-
-    fn free_part(&self, _ptr: *mut u8, _oldsize: usize, _newsize: usize) -> bool {
-        false
-    }
-
-    fn free(&self, _ptr: *mut u8, _size: usize) -> bool {
-        return false;
-    }
-
-    fn can_release_part(&self, _flags: u32) -> bool {
-        false
-    }
-
-    fn allocates_zeros(&self) -> bool {
-        false
-    }
-
-    fn page_size(&self) -> usize {
-        0x1000
-    }
-}
+use crate::{alloc::{self, GlobalAlloc, System}};
+use snmalloc_edp::*;
 
 #[stable(feature = "alloc_system_type", since = "1.28.0")]
 unsafe impl GlobalAlloc for System {
     #[inline]
-    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        // SAFETY: the caller must uphold the safety contract for `malloc`
-        unsafe { DLMALLOC.lock().malloc(layout.size(), layout.align()) }
+    unsafe fn alloc(&self, layout: alloc::Layout) -> *mut u8 {
+        unsafe { sn_rust_alloc(layout.align(), layout.size()) }
     }
 
     #[inline]
-    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+    unsafe fn alloc_zeroed(&self, layout: alloc::Layout) -> *mut u8 {
         // SAFETY: the caller must uphold the safety contract for `malloc`
-        unsafe { DLMALLOC.lock().calloc(layout.size(), layout.align()) }
+        unsafe { sn_rust_alloc_zeroed(layout.align(), layout.size()) }
     }
 
     #[inline]
-    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: alloc::Layout) {
         // SAFETY: the caller must uphold the safety contract for `malloc`
-        unsafe { DLMALLOC.lock().free(ptr, layout.size(), layout.align()) }
+        unsafe { sn_rust_dealloc(ptr, layout.align(), layout.size()) }
     }
 
     #[inline]
-    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+    unsafe fn realloc(&self, ptr: *mut u8, layout: alloc::Layout, new_size: usize) -> *mut u8 {
         // SAFETY: the caller must uphold the safety contract for `malloc`
-        unsafe { DLMALLOC.lock().realloc(ptr, layout.size(), layout.align(), new_size) }
+        unsafe { sn_rust_realloc(ptr, layout.align(), layout.size(), new_size) }
     }
 }
 
@@ -88,11 +32,17 @@ unsafe impl GlobalAlloc for System {
 #[cfg(not(test))]
 #[no_mangle]
 pub unsafe extern "C" fn __rust_c_alloc(size: usize, align: usize) -> *mut u8 {
-    unsafe { crate::alloc::alloc(Layout::from_size_align_unchecked(size, align)) }
+    unsafe { crate::alloc::alloc(crate::alloc::Layout::from_size_align_unchecked(size, align)) }
 }
 
 #[cfg(not(test))]
 #[no_mangle]
 pub unsafe extern "C" fn __rust_c_dealloc(ptr: *mut u8, size: usize, align: usize) {
-    unsafe { crate::alloc::dealloc(ptr, Layout::from_size_align_unchecked(size, align)) }
+    unsafe { crate::alloc::dealloc(ptr, crate::alloc::Layout::from_size_align_unchecked(size, align)) }
+}
+
+#[cfg(not(test))]
+#[no_mangle]
+pub extern "C" fn __rust_get_thread_allocator() -> *mut Alloc {
+    unsafe{ crate::sys::abi::tls::get_tls_ptr(crate::sys::abi::tls::TlsIndex::AllocPtr) as *mut Alloc }
 }