[libc-commits] [libc] [libc] make integration test malloc work properly when threaded (PR #151622)

[via libc-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-libc

Author: Schrodinger ZHU Yifan (SchrodingerZhu)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/151622.diff


2 Files Affected:

- (modified) libc/test/IntegrationTest/CMakeLists.txt (+1) 
- (modified) libc/test/IntegrationTest/test.cpp (+16-4) 


``````````diff
diff --git a/libc/test/IntegrationTest/CMakeLists.txt b/libc/test/IntegrationTest/CMakeLists.txt
index 4a999407d48d7..1763ed8b932fe 100644
--- a/libc/test/IntegrationTest/CMakeLists.txt
+++ b/libc/test/IntegrationTest/CMakeLists.txt
@@ -12,5 +12,6 @@ add_object_library(
     test.h
   DEPENDS
     libc.src.__support.OSUtil.osutil
+    libc.src.__support.CPP.atomic
     ${arch_specific_deps}
 )
diff --git a/libc/test/IntegrationTest/test.cpp b/libc/test/IntegrationTest/test.cpp
index 871bdf0dc562b..1a489403f37b5 100644
--- a/libc/test/IntegrationTest/test.cpp
+++ b/libc/test/IntegrationTest/test.cpp
@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "src/__support/CPP/atomic.h"
 #include "src/__support/common.h"
 #include "src/__support/macros/config.h"
 #include <stddef.h>
@@ -65,14 +66,25 @@ int atexit(void (*func)(void)) { return LIBC_NAMESPACE::atexit(func); }
 
 static constexpr uint64_t MEMORY_SIZE = 16384;
 static uint8_t memory[MEMORY_SIZE];
-static uint8_t *ptr = memory;
+static LIBC_NAMESPACE::cpp::Atomic<size_t> used = 0;
 
 extern "C" {
 
 void *malloc(size_t s) {
-  void *mem = ptr;
-  ptr += s;
-  return static_cast<uint64_t>(ptr - memory) >= MEMORY_SIZE ? nullptr : mem;
+  constexpr size_t MAX_ALIGNMENT = alignof(long double);
+  if (s > MEMORY_SIZE)
+    return nullptr; // Not enough memory.
+  s = s + MAX_ALIGNMENT -
+      (s % MAX_ALIGNMENT); // Align the size to the max alignment.
+  for (;;) {
+    size_t current_used = used.load();
+    if (current_used + s > MEMORY_SIZE)
+      return nullptr; // Not enough memory.
+
+    // Try to reserve the memory.
+    if (used.compare_exchange_strong(current_used, current_used + s))
+      return &memory[current_used];
+  }
 }
 
 void free(void *) {}

``````````

</details>


https://github.com/llvm/llvm-project/pull/151622