From eb9cc253cb048b6dbf2fcd73ac55b5eda0184ed3 Mon Sep 17 00:00:00 2001
From: Siva Chandra Reddy <sivachandra@google.com>
Date: Wed, 21 Dec 2022 07:27:08 +0000
Subject: [PATCH] [libc] Gracefully handle allocation failures around
 BlockStore.

Reviewed By: lntue

Differential Revision: https://reviews.llvm.org/D140459
---
 libc/src/__support/CMakeLists.txt           |  2 ++
 libc/src/__support/blockstore.h             | 19 +++++++++++++------
 libc/src/stdlib/atexit.cpp                  | 10 +++-------
 libc/test/src/__support/blockstore_test.cpp | 10 ++++++----
 4 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/libc/src/__support/CMakeLists.txt b/libc/src/__support/CMakeLists.txt
index 3ca589b..fc80c53 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 44c8b98..43b6483 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 @@ public:
 
   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 @@ public:
     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 @@ public:
       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 7725463..88a7d1a 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 3edbb92..f844999 100644
--- a/libc/test/src/__support/blockstore_test.cpp
+++ b/libc/test/src/__support/blockstore_test.cpp
@@ -21,7 +21,7 @@ public:
   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 @@ public:
     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 @@ public:
     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);
   }
 };
-- 
2.7.4