Fix and improve few things

libcoz/CMakeLists.txt

@@ -1,5 +1,6 @@
 set(sources
     ${PROJECT_SOURCE_DIR}/include/coz.h
+    alloc_shims.h
     inspect.cpp
     inspect.h
     libcoz.cpp

libcoz/alloc_shims.cpp

@@ -17,23 +17,24 @@ static spinlock init_lock;
 static spinlock mode_lock;
 static std::atomic<pid_t> thread_using_shim{0};
 
-static constexpr size_t memory_pool_size = 1000 * sizeof(std::max_align_t);
+static constexpr size_t memory_pool_size = 1000 * alignof(std::max_align_t);
 alignas(std::max_align_t) static char memory_pool[memory_pool_size];
 
-static void init_shims_if_needed();
+static void lazy_init();
 
 static void* first_malloc(size_t size) {
-  init_shims_if_needed();
-  // Use what ever malloc is set after init.
+  lazy_init();
   return malloc(size);
 }
 
 static void* first_calloc(size_t nmemb, size_t size) {
-  init_shims_if_needed();
-  // Use what ever calloc is set after init.
+  lazy_init();
   return calloc(nmemb, size);
 }
 
+static void* (*in_use_malloc)(size_t size) = first_malloc;
+static void* (*in_use_calloc)(size_t nmemb, size_t size) = first_calloc;
+
 static void* (*real_malloc)(size_t size) = nullptr;
 static void (*real_free)(void* ptr) = nullptr;
 static void* (*real_calloc)(size_t nmemb, size_t size) = nullptr;
@@ -61,10 +62,10 @@ static void* dummy_malloc(size_t size) {
 
   static char* first_unallocated{memory_pool};
 
-  // Make size multiple of sizeof(max_align_t), to keep addresses aligned.
+  // Make size multiple of alignof(max_align_t), to keep addresses aligned.
   constexpr std::uintmax_t all_ones = ~(std::uintmax_t{});
-  size = (size + sizeof(std::max_align_t) - 1) &
-      (all_ones * sizeof(std::max_align_t)); // this is same as shifting left.
+  size = (size + alignof(std::max_align_t) - 1) &
+      (all_ones * alignof(std::max_align_t)); // this is same as shifting left.
 
   char* result = first_unallocated;
   first_unallocated += size;
@@ -81,18 +82,15 @@ static void* dummy_calloc(size_t nmemb, size_t size) {
   return ptr;
 }
 
-static void* (*in_use_malloc)(size_t size) = first_malloc;
-static void* (*in_use_calloc)(size_t nmemb, size_t size) = first_calloc;
-
-static void do_set_dummy_allocs() {
+static void set_dummy_allocs_impl() {
   // If another thread ends up in dummy_malloc, make sure it knows it's in the
   // wrong place.
   thread_using_shim.store(gettid());
   in_use_malloc = dummy_malloc;
   in_use_calloc = dummy_calloc;
 }
 
-static void do_restore_real_allocs() {
+static void restore_real_allocs_impl() {
   in_use_malloc = real_malloc;
   in_use_calloc = real_calloc;
   thread_using_shim.store(0);
@@ -120,46 +118,36 @@ static void find_real_functions() {
   }
 }
 
-static void init_shims_if_needed() {
-  // It is important to note that this function can be called on first try to
-  // allocated memory (first call to malloc or calloc, delegated to
-  // first_malloc or first_calloc), or first call to
-  // coz_lock_and_set_dummy_alloc_shims.
-  // This is because, either of these two events requires allocation,
-  // and everything becomes much simpler, if we initialize everything in advance
-  // in both cases.
-
+static void lazy_init() {
   // First check to improve performance, avoid locking if unnecessary.
-  if (initialized.load())
+  if (initialized.load(std::memory_order_acquire))
     return;
 
   init_lock.lock();
 
   // Another check, this time to make sure no one acquired the lock between
   // first check and our attempt to acquire it.
-  if (initialized.load())
+  if (initialized.load(std::memory_order_relaxed))
     return;
 
   // Allocations could be made by libdl while we search for real functions.
   // Prepare dummy allocs for this (first_malloc is no longer needed, we are
   // already initializing).
-  do_set_dummy_allocs();
+  set_dummy_allocs_impl();
 
   // Now find real functions.
   find_real_functions();
 
-  // Now that we have real functions, use them. If we are initialized from,
-  // coz_lock_and_set_dummy_alloc_shims, it will rewrite it to dummy functions
-  // again, but it doesn't matter.
-  do_restore_real_allocs();
+  // Now that we have real functions, use them by default.
+  restore_real_allocs_impl();
 
-  initialized.store(true);
+  initialized.store(true, std::memory_order_release);
   init_lock.unlock();
 }
 
 extern "C" {
 void coz_lock_and_set_dummy_alloc_shims() {
-  init_shims_if_needed();
+  lazy_init();
 
   // This is to make sure, only one thread resolves real symbols at the time.
   // That in turn makes it possible, to ensure dummy malloc is called only
@@ -168,11 +156,11 @@ void coz_lock_and_set_dummy_alloc_shims() {
   // resolved from one thread, and each resolve_* function is called only once,
   // during runtime of the program.
   mode_lock.lock();
-  do_set_dummy_allocs();
+  set_dummy_allocs_impl();
 }
 
 void coz_restore_real_alloc_shims_and_unlock() {
-  do_restore_real_allocs();
+  restore_real_allocs_impl();
   mode_lock.unlock();
 }