[libc-commits] [libc] eb9cc25 - [libc] Gracefully handle allocation failures around BlockStore.

[Siva Chandra Reddy via libc-commits]

Author: Siva Chandra Reddy
Date: 2022-12-21T21:05:09Z
New Revision: eb9cc253cb048b6dbf2fcd73ac55b5eda0184ed3

URL: https://github.com/llvm/llvm-project/commit/eb9cc253cb048b6dbf2fcd73ac55b5eda0184ed3
DIFF: https://github.com/llvm/llvm-project/commit/eb9cc253cb048b6dbf2fcd73ac55b5eda0184ed3.diff

LOG: [libc] Gracefully handle allocation failures around BlockStore.

Reviewed By: lntue

Differential Revision: https://reviews.llvm.org/D140459

Added: 
    

Modified: 
    libc/src/__support/CMakeLists.txt
    libc/src/__support/blockstore.h
    libc/src/stdlib/atexit.cpp
    libc/test/src/__support/blockstore_test.cpp

Removed: 
    


################################################################################
diff  --git a/libc/src/__support/CMakeLists.txt b/libc/src/__support/CMakeLists.txt
index 3ca589b71781e..fc80c5389ddfc 100644
--- a/libc/src/__support/CMakeLists.txt
+++ b/libc/src/__support/CMakeLists.txt
@@ -4,6 +4,8 @@ add_header_library(
   blockstore
   HDRS
     blockstore.h
+  DEPENDS
+    libc.src.__support.CPP.new
 )
 
 add_header_library(

diff  --git a/libc/src/__support/blockstore.h b/libc/src/__support/blockstore.h
index 44c8b98551d93..43b6483846665 100644
--- a/libc/src/__support/blockstore.h
+++ b/libc/src/__support/blockstore.h
@@ -9,9 +9,10 @@
 #ifndef LLVM_LIBC_SUPPORT_BLOCKSTORE_H
 #define LLVM_LIBC_SUPPORT_BLOCKSTORE_H
 
+#include <src/__support/CPP/new.h>
+
 #include <stddef.h>
 #include <stdint.h>
-#include <stdlib.h>
 
 // TODO: fix our assert.h to make it useable
 #define assert(x)                                                              \
@@ -114,7 +115,10 @@ class BlockStore {
 
   T *new_obj() {
     if (fill_count == BLOCK_SIZE) {
-      auto new_block = reinterpret_cast<Block *>(::malloc(sizeof(Block)));
+      AllocChecker ac;
+      auto new_block = new (ac) Block();
+      if (!ac)
+        return nullptr;
       if (REVERSE_ORDER) {
         new_block->next = current;
       } else {
@@ -129,9 +133,12 @@ class BlockStore {
     return obj;
   }
 
-  void push_back(const T &value) {
+  [[nodiscard]] bool push_back(const T &value) {
     T *ptr = new_obj();
+    if (ptr == nullptr)
+      return false;
     *ptr = value;
+    return true;
   }
 
   T &back() {
@@ -153,7 +160,7 @@ class BlockStore {
       current->next = nullptr;
     }
     if (last != &first)
-      ::free(last);
+      delete last;
     fill_count = BLOCK_SIZE;
   }
 
@@ -182,14 +189,14 @@ void BlockStore<T, BLOCK_SIZE, REVERSE_ORDER>::destroy(
     while (current->next != nullptr) {
       auto temp = current;
       current = current->next;
-      free(temp);
+      delete temp;
     }
   } else {
     auto current = block_store->first.next;
     while (current != nullptr) {
       auto temp = current;
       current = current->next;
-      free(temp);
+      delete temp;
     }
   }
   block_store->current = nullptr;

diff  --git a/libc/src/stdlib/atexit.cpp b/libc/src/stdlib/atexit.cpp
index 7725463b8eb74..88a7d1a579a0d 100644
--- a/libc/src/stdlib/atexit.cpp
+++ b/libc/src/stdlib/atexit.cpp
@@ -67,13 +67,9 @@ void call_exit_callbacks() {
 } // namespace internal
 
 static int add_atexit_unit(const AtExitUnit &unit) {
-  // TODO: Use the return value of push_back and bubble it to the public
-  // function as error return value. Note that a BlockStore push_back can
-  // fail because of allocation failure. Likewise, a FixedVector push_back
-  // can fail when it is full.
-  handler_list_mtx.lock();
-  exit_callbacks.push_back(unit);
-  handler_list_mtx.unlock();
+  MutexLock lock(&handler_list_mtx);
+  if (!exit_callbacks.push_back(unit))
+    return -1;
   return 0;
 }
 

diff  --git a/libc/test/src/__support/blockstore_test.cpp b/libc/test/src/__support/blockstore_test.cpp
index 3edbb920144bd..f844999b06e99 100644
--- a/libc/test/src/__support/blockstore_test.cpp
+++ b/libc/test/src/__support/blockstore_test.cpp
@@ -21,7 +21,7 @@ class LlvmLibcBlockStoreTest : public __llvm_libc::testing::Test {
   void populate_and_iterate() {
     __llvm_libc::cpp::BlockStore<Element, BLOCK_SIZE, REVERSE> block_store;
     for (int i = 0; i < int(ELEMENT_COUNT); ++i)
-      block_store.push_back({i, 2 * i, 3 * unsigned(i)});
+      ASSERT_TRUE(block_store.push_back({i, 2 * i, 3 * unsigned(i)}));
     auto end = block_store.end();
     int i = 0;
     for (auto iter = block_store.begin(); iter != end; ++iter, ++i) {
@@ -46,7 +46,7 @@ class LlvmLibcBlockStoreTest : public __llvm_libc::testing::Test {
     using __llvm_libc::cpp::BlockStore;
     BlockStore<int, 4, REVERSE> block_store;
     for (int i = 0; i < 20; i++)
-      block_store.push_back(i);
+      ASSERT_TRUE(block_store.push_back(i));
     for (int i = 19; i >= 0; i--, block_store.pop_back())
       ASSERT_EQ(block_store.back(), i);
     block_store.destroy(&block_store);
@@ -57,9 +57,11 @@ class LlvmLibcBlockStoreTest : public __llvm_libc::testing::Test {
     BlockStore<int, 2, REVERSE> block_store;
 
     ASSERT_TRUE(block_store.empty());
-    block_store.push_back(1);
-    for (int i = 0; i < 10; i++, block_store.push_back(1))
+    ASSERT_TRUE(block_store.push_back(1));
+    for (int i = 0; i < 10; i++) {
       ASSERT_FALSE(block_store.empty());
+      ASSERT_TRUE(block_store.push_back(1));
+    }
     block_store.destroy(&block_store);
   }
 };