From 8612d1228a76c306898e784fe0a1829ab67d4cd2 Mon Sep 17 00:00:00 2001
From: Daan <daanl@outlook.com>
Date: Wed, 15 Dec 2021 18:04:02 -0800
Subject: [PATCH] wip: initial fix for race condition on sbrk/wasm_memory_grow;
 issue #497

---
 src/os.c | 87 +++++++++++++++++++++++++++++++++++---------------------
 1 file changed, 55 insertions(+), 32 deletions(-)

diff --git a/src/os.c b/src/os.c
index 603ec04a..dc2afc78 100644
--- a/src/os.c
+++ b/src/os.c
@@ -375,42 +375,69 @@ static void* mi_win_virtual_alloc(void* addr, size_t size, size_t try_alignment,
   return p;
 }
 
-#elif defined(MI_USE_SBRK)
+#elif defined(MI_USE_SBRK) || defined(__wasi__)
+
+#if defined(MI_USE_SBRK) 
 #define MI_SBRK_FAIL ((void*)(-1)) 
-static void* mi_sbrk_heap_grow(size_t size, size_t try_alignment) {
-  void* pbase0 = sbrk(0);
-  if (pbase0 == MI_SBRK_FAIL) {
+static void* mi_memory_grow( size_t size ) {
+  void* p = sbrk(size);
+  if (p == MI_SBRK_FAIL) {
     _mi_warning_message("unable to allocate sbrk() OS memory (%zu bytes)\n", size);
     errno = ENOMEM;
     return NULL;
   }
-  uintptr_t base = (uintptr_t)pbase0;
-  uintptr_t aligned_base = _mi_align_up(base, (uintptr_t) try_alignment);
-  size_t alloc_size = _mi_align_up( aligned_base - base + size, _mi_os_page_size());
-  mi_assert(alloc_size >= size && (alloc_size % _mi_os_page_size()) == 0);
-  if (alloc_size < size) return NULL;
-  void* pbase1 = sbrk(alloc_size);
-  if (pbase1 == MI_SBRK_FAIL) {
-    _mi_warning_message("unable to allocate sbrk() OS memory (%zu bytes, %zu requested)\n", size, alloc_size);
-    errno = ENOMEM;
-    return NULL;
-  }
-  mi_assert(pbase0 == pbase1);
-  return (void*)aligned_base;
+  return p;
 }
+#elif defined(__wasi__)
+static void* mi_memory_grow( size_t size ) {
+  if (size > 0) {
+    size_t base = __builtin_wasm_memory_grow( 0, _mi_divide_up(size, _mi_os_page_size()) );
+    if (base == SIZE_MAX) {
+      _mi_warning_message("unable to allocate wasm_memory_grow OS memory (%zu bytes)\n", size);    
+      errno = ENOMEM;
+      return NULL;
+    }
+    return (void*)(base * _mi_os_page_size());
+  }
+  else {
+    size_t base = __builtin_wasm_memory_size(0);
+    if (base == SIZE_MAX) {
+      errno = ENOMEM;      
+      return NULL; 
+    }
+    return (void*)(base * _mi_os_page_size());
+  }
+}
+#endif
 
-#elif defined(__wasi__)  
- // currently unused as we use sbrk() on wasm
-static void* mi_wasm_heap_grow(size_t size, size_t try_alignment) {
-  uintptr_t base = __builtin_wasm_memory_size(0) * _mi_os_page_size();
-  uintptr_t aligned_base = _mi_align_up(base, (uintptr_t) try_alignment);
+static void* mi_heap_grow(size_t size, size_t try_alignment) {
+  if (try_alignment == 0) { try_alignment = _mi_os_page_size(); };
+  void* pbase0 = mi_memory_grow(0);
+  if (pbase0 == NULL) { return NULL; }
+  uintptr_t base = (uintptr_t)pbase0;
+  uintptr_t aligned_base = _mi_align_up(base, try_alignment);
   size_t alloc_size = _mi_align_up( aligned_base - base + size, _mi_os_page_size());
   mi_assert(alloc_size >= size && (alloc_size % _mi_os_page_size()) == 0);
   if (alloc_size < size) return NULL;
-  if (__builtin_wasm_memory_grow(0, alloc_size / _mi_os_page_size()) == SIZE_MAX) {
-    _mi_warning_message("unable to allocate wasm_memory_grow() OS memory (%zu bytes, %zu requested)\n", size, alloc_size);
-    errno = ENOMEM;
-    return NULL;
+  void* pbase1 = mi_memory_grow(alloc_size);
+  if (pbase1 == NULL) { return NULL; }
+  if (pbase0 != pbase1) {
+    // another thread allocated in-between; now we may not be able to align correctly
+    base = (uintptr_t)pbase1;
+    aligned_base = _mi_align_up(base, try_alignment);
+    if (aligned_base + size > base + alloc_size) {
+      // we do not have enough space after alignment; since we cannot shrink safely,
+      // we waste the space :-( and allocate fresh with guaranteed enough overallocation
+      alloc_size = _mi_align_up( size + try_alignment, _mi_os_page_size() );
+      errno = 0;
+      void* pbase2 = mi_memory_grow( alloc_size );
+      if (pbase2 == NULL) { return NULL; }
+      aligned_base = _mi_align_up(base, try_alignment);
+      mi_assert_internal(aligned_base + size <= (uintptr_t)pbase2 + alloc_size);
+    }    
+  }
+  else {
+    mi_assert_internal(aligned_base + size <= (uintptr_t)pbase1 + alloc_size);
   }
   return (void*)aligned_base;
 }
@@ -637,14 +664,10 @@ static void* mi_os_mem_alloc(size_t size, size_t try_alignment, bool commit, boo
     int flags = MEM_RESERVE;
     if (commit) flags |= MEM_COMMIT;
     p = mi_win_virtual_alloc(NULL, size, try_alignment, flags, false, allow_large, is_large);
-  #elif defined(MI_USE_SBRK)
+  #elif defined(MI_USE_SBRK) || defined(__wasi__)
     MI_UNUSED(allow_large);
     *is_large = false;
-    p = mi_sbrk_heap_grow(size, try_alignment);
-  #elif defined(__wasi__)
-    MI_UNUSED(allow_large);
-    *is_large = false;
-    p = mi_wasm_heap_grow(size, try_alignment);
+    p = mi_heap_grow(size, try_alignment);
   #else
     int protect_flags = (commit ? (PROT_WRITE | PROT_READ) : PROT_NONE);
     p = mi_unix_mmap(NULL, size, try_alignment, protect_flags, false, allow_large, is_large);