update lock primitive; fix arena exclusive allocation

include/mimalloc/atomic.h

@@ -1,5 +1,5 @@
 /* ----------------------------------------------------------------------------
-Copyright (c) 2018-2023 Microsoft Research, Daan Leijen
+Copyright (c) 2018-2024 Microsoft Research, Daan Leijen
 This is free software; you can redistribute it and/or modify it under the
 terms of the MIT license. A copy of the license can be found in the file
 "LICENSE" at the root of this distribution.
@@ -411,8 +411,11 @@ static inline void mi_atomic_yield(void) {
 #pragma warning(disable:26110)  // unlock with holding lock
 #endif
 
+#define mi_lock(lock)    for(bool _go = (mi_lock_acquire(lock),true); _go; (mi_lock_release(lock), _go=false) )
+
 #if defined(_WIN32)
 
+#if 1
 #define mi_lock_t  SRWLOCK   // slim reader-writer lock
 
 static inline bool mi_lock_try_acquire(mi_lock_t* lock) {
@@ -432,6 +435,30 @@ static inline void mi_lock_done(mi_lock_t* lock) {
   // nothing
 }
 
+#else
+#define mi_lock_t  CRITICAL_SECTION
+
+static inline bool mi_lock_try_acquire(mi_lock_t* lock) {
+  return TryEnterCriticalSection(lock);
+  
+}
+static inline void mi_lock_acquire(mi_lock_t* lock) {
+  EnterCriticalSection(lock);
+  
+}
+static inline void mi_lock_release(mi_lock_t* lock) {
+  LeaveCriticalSection(lock);
+  
+}
+static inline void mi_lock_init(mi_lock_t* lock) {
+  InitializeCriticalSection(lock);
+  
+}
+static inline void mi_lock_done(mi_lock_t* lock) {
+  DeleteCriticalSection(lock);
+  
+}
+#endif
 
 #elif defined(MI_USE_PTHREADS)
 
@@ -506,6 +533,4 @@ static inline void mi_lock_done(mi_lock_t* lock) {
 #endif
 
 
-
-
 #endif // __MIMALLOC_ATOMIC_H

src/arena-abandon.c

@@ -120,11 +120,7 @@ static void mi_arena_segment_os_mark_abandoned(mi_segment_t* segment) {
   mi_assert(segment->memid.memkind != MI_MEM_ARENA);
   // not in an arena; we use a list of abandoned segments
   mi_subproc_t* const subproc = segment->subproc;
-  if (!mi_lock_acquire(&subproc->abandoned_os_lock)) {
-    _mi_error_message(EFAULT, "internal error: failed to acquire the abandoned (os) segment lock to mark abandonment");
-    // we can continue but cannot visit/reclaim such blocks..
-  }
-  else {
+  mi_lock(&subproc->abandoned_os_lock) {
     // push on the tail of the list (important for the visitor)
     mi_segment_t* prev = subproc->abandoned_os_list_tail;
     mi_assert_internal(prev == NULL || prev->abandoned_os_next == NULL);
@@ -138,7 +134,6 @@ static void mi_arena_segment_os_mark_abandoned(mi_segment_t* segment) {
     mi_atomic_increment_relaxed(&subproc->abandoned_os_list_count);
     mi_atomic_increment_relaxed(&subproc->abandoned_count);
     // and release the lock
-    mi_lock_release(&subproc->abandoned_os_lock);
   }
   return;
 }
@@ -251,7 +246,7 @@ static mi_segment_t* mi_arena_segment_clear_abandoned_next_field(mi_arena_field_
         if mi_unlikely(field != 0) { // skip zero fields quickly
           // we only take the arena lock if there are actually abandoned segments present
           if (!has_lock && mi_option_is_enabled(mi_option_visit_abandoned)) {
-            has_lock = (previous->visit_all ? mi_lock_acquire(&arena->abandoned_visit_lock) : mi_lock_try_acquire(&arena->abandoned_visit_lock));
+            has_lock = (previous->visit_all ? (mi_lock_acquire(&arena->abandoned_visit_lock),true) : mi_lock_try_acquire(&arena->abandoned_visit_lock));
             if (!has_lock) {
               if (previous->visit_all) {
                 _mi_error_message(EFAULT, "internal error: failed to visit all abandoned segments due to failure to acquire the visitor lock");
@@ -289,7 +284,7 @@ static mi_segment_t* mi_arena_segment_clear_abandoned_next_list(mi_arena_field_c
   // we only allow one thread per sub-process to do to visit guarded by the `abandoned_os_visit_lock`.
   // The lock is released when the cursor is released.
   if (!previous->hold_visit_lock) {
-    previous->hold_visit_lock = (previous->visit_all ? mi_lock_acquire(&previous->subproc->abandoned_os_visit_lock)
+    previous->hold_visit_lock = (previous->visit_all ? (mi_lock_acquire(&previous->subproc->abandoned_os_visit_lock),true)
                                                      : mi_lock_try_acquire(&previous->subproc->abandoned_os_visit_lock));
     if (!previous->hold_visit_lock) {
       if (previous->visit_all) {
@@ -301,8 +296,7 @@ static mi_segment_t* mi_arena_segment_clear_abandoned_next_list(mi_arena_field_c
   // One list entry at a time
   while (previous->os_list_count > 0) {
     previous->os_list_count--;
-    const bool has_lock = mi_lock_acquire(&previous->subproc->abandoned_os_lock); // this could contend with concurrent OS block abandonment and reclaim from `free`
-    if (has_lock) {
+    mi_lock_acquire(&previous->subproc->abandoned_os_lock); // this could contend with concurrent OS block abandonment and reclaim from `free`
     mi_segment_t* segment = previous->subproc->abandoned_os_list;
     // pop from head of the list, a subsequent mark will push at the end (and thus we iterate through os_list_count entries)
     if (segment == NULL || mi_arena_segment_os_clear_abandoned(segment, false /* we already have the lock */)) {
@@ -312,11 +306,6 @@ static mi_segment_t* mi_arena_segment_clear_abandoned_next_list(mi_arena_field_c
     // already abandoned, try again
     mi_lock_release(&previous->subproc->abandoned_os_lock);
   }
-    else {
-      _mi_error_message(EFAULT, "failed to acquire abandoned OS list lock during abandoned block visit\n");
-      return NULL;
-    }
-  }
   // done
   mi_assert_internal(previous->os_list_count == 0);
   return NULL;

src/arena.c

@@ -394,8 +394,9 @@ void* _mi_arena_alloc_aligned(size_t size, size_t alignment, size_t align_offset
   const int numa_node = _mi_os_numa_node(); // current numa node
 
   // try to allocate in an arena if the alignment is small enough and the object is not too small (as for heap meta data)
-  if (!mi_option_is_enabled(mi_option_disallow_arena_alloc) || req_arena_id != _mi_arena_id_none()) {  // is arena allocation allowed?
-    if (size >= MI_ARENA_MIN_OBJ_SIZE && alignment <= MI_SEGMENT_ALIGN && align_offset == 0) {
+  if (!mi_option_is_enabled(mi_option_disallow_arena_alloc)) {  // is arena allocation allowed?
+    if (size >= MI_ARENA_MIN_OBJ_SIZE && alignment <= MI_SEGMENT_ALIGN && align_offset == 0) 
+    {
       void* p = mi_arena_try_alloc(numa_node, size, alignment, commit, allow_large, req_arena_id, memid);
       if (p != NULL) return p;
 

src/init.c

@@ -257,11 +257,10 @@ void mi_subproc_delete(mi_subproc_id_t subproc_id) {
   mi_subproc_t* subproc = _mi_subproc_from_id(subproc_id);
   // check if there are no abandoned segments still..
   bool safe_to_delete = false;
-  if (mi_lock_acquire(&subproc->abandoned_os_lock)) {
+  mi_lock(&subproc->abandoned_os_lock) {
     if (subproc->abandoned_os_list == NULL) {
       safe_to_delete = true;
     }
-    mi_lock_release(&subproc->abandoned_os_lock);
   }
   if (!safe_to_delete) return;
   // safe to release