src/lib-alloc.c: refactor C init stubs

Inria-Prosecco · Oct 19, 2023 · f0a5f5a · f0a5f5a
1 parent db25dda
commit f0a5f5a
Showing 1 changed file with 58 additions and 64 deletions.
diff --git a/src/lib-alloc.c b/src/lib-alloc.c
@@ -5,14 +5,13 @@
 #include "krmlinit.h"
 #include "Config.h"
 #include "StarMalloc.h"
+#include "internal/StarMalloc.h"
 
-#define N_ARENA 4
-
-static uint32_t init_status = 0UL;
-static pthread_mutex_t m = PTHREAD_MUTEX_INITIALIZER;
+//TODO: remove, mark Config as const (correctness + performance)
+#define N_ARENAS 4
 
 __attribute__((tls_model("initial-exec")))
-static _Thread_local unsigned thread_arena = N_ARENA;
+static _Thread_local unsigned thread_arena = N_ARENAS;
 static atomic_uint thread_arena_counter = 0;
 
 uint8_t* StarMalloc_memcpy_u8(uint8_t* dest, uint8_t* src, size_t n) {
@@ -23,81 +22,76 @@ uint8_t* StarMalloc_memset_u8(uint8_t* dest, uint8_t v, size_t n) {
   return (uint8_t*) memset((void*) dest, v, n);
 }
 
-void* malloc(size_t size) {
-  if (! init_status) {
-    pthread_mutex_lock(&m);
-    krmlinit_globals();
-    init_status=1UL;
+void* _Atomic slab_region_ptr;
+
+// hm-alike initialization
+static inline bool is_init(void) {
+  void* ptr = atomic_load_explicit(&slab_region_ptr, memory_order_acquire);
+  return ptr != NULL;
+}
+
+// hm-alike
+static void init_slow_path(void) {
+  static pthread_mutex_t m = PTHREAD_MUTEX_INITIALIZER;
+  pthread_mutex_lock(&m);
+  if (is_init()) {
     pthread_mutex_unlock(&m);
+    return;
   }
-  if (thread_arena == N_ARENA) {
-    thread_arena = thread_arena_counter++ % N_ARENA;
-  }
-  return StarMalloc_malloc(thread_arena, size);
+  krmlinit_globals();
+  atomic_store_explicit(&slab_region_ptr, Main_Meta_sc_all.slab_region, memory_order_release);
+  pthread_mutex_unlock(&m);
+  // TODO: pthread_atfork
 }
 
-void* aligned_alloc(size_t alignment, size_t size) {
-  if (! init_status) {
-    pthread_mutex_lock(&m);
-    krmlinit_globals();
-    init_status=1UL;
-    pthread_mutex_unlock(&m);
+// hm-alike
+static inline unsigned init(void) {
+  unsigned arena = thread_arena;
+  if (arena < N_ARENAS) {
+    return arena;
   }
-  if (thread_arena == N_ARENA) {
-    thread_arena = thread_arena_counter++ % N_ARENA;
+  thread_arena = arena = thread_arena_counter++ % N_ARENAS;
+  if (!is_init ()) {
+    init_slow_path();
   }
-  return StarMalloc_aligned_alloc(thread_arena, alignment, size);
+  return arena;
 }
 
-void free(void *ptr) {
-  if (! init_status) {
-    pthread_mutex_lock(&m);
-    krmlinit_globals();
-    init_status=1UL;
-    pthread_mutex_unlock(&m);
-  }
-  //printf("free ptr: %p\n", ptr);
-  bool b = StarMalloc_free(ptr);
-  //printf("  result: %b\n");
-  return;
+void* malloc(size_t size) {
+  unsigned arena = init();
+  return StarMalloc_malloc(arena, size);
 }
 
-void* realloc(void* ptr, size_t new_size) {
-  if (! init_status) {
-    pthread_mutex_lock(&m);
-    krmlinit_globals();
-    init_status=1UL;
-    pthread_mutex_unlock(&m);
-  }
-  if (thread_arena == N_ARENA) {
-    thread_arena = thread_arena_counter++ % N_ARENA;
-  }
-  //printf("ptr: %p, new_size: %lu\n", ptr, new_size);
-  void* new_ptr = StarMalloc_realloc(thread_arena, ptr, new_size);
-  //printf("  new_ptr: %p\n", new_ptr);
-  return new_ptr;
+void* aligned_alloc(size_t alignment, size_t size) {
+  unsigned arena = init();
+  return StarMalloc_aligned_alloc(arena, alignment, size);
 }
 
 void* calloc(size_t nb_elem, size_t size_elem) {
-  if (! init_status) {
-    pthread_mutex_lock(&m);
-    krmlinit_globals();
-    init_status=1UL;
-    pthread_mutex_unlock(&m);
-  }
-  if (thread_arena == N_ARENA) {
-    thread_arena = thread_arena_counter++ % N_ARENA;
-  }
-  uint8_t* ptr = StarMalloc_calloc(thread_arena, nb_elem, size_elem);
-  return (void*) ptr;
+  unsigned arena = init();
+  return StarMalloc_calloc(arena, nb_elem, size_elem);
+}
+
+void* realloc(void* ptr, size_t new_size) {
+  unsigned arena = init();
+  return StarMalloc_realloc(arena, ptr, new_size);
+}
+
+void free(void *ptr) {
+  unsigned arena = init();
+  //printf("free ptr: %p\n", ptr);
+  bool b = StarMalloc_free(ptr);
+  assert (b);
+  //printf("  result: %b\n");
+  return;
 }
 
 size_t malloc_usable_size(void* ptr) {
-  if (! init_status) {
-    pthread_mutex_lock(&m);
-    krmlinit_globals();
-    init_status=1UL;
-    pthread_mutex_unlock(&m);
+  // not needed for correctness, but way faster
+  if (ptr == NULL) {
+    return 0;
   }
+  // use enforce_init
+  init();
   return StarMalloc_getsize(ptr);
 }