[libc] Gracefully handle allocation failures around BlockStore.

libc/src/__support/CMakeLists.txt

 add_subdirectory(CPP)
 
 add_header_library(
   blockstore
   HDRS
     blockstore.h
+  DEPENDS
+    libc.src.__support.CPP.new
 )
 
 add_header_library(
   named_pair
   HDRS
     named_pair.h
 )
 

libc/src/__support/blockstore.h

 //===-- A data structure which stores data in blocks  -----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #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)                                                              \
   if (!(x))                                                                    \
   __builtin_trap()
 
 namespace __llvm_libc {
 namespace cpp {
bool operator!=(const iterator &rhs) const {
       return block != rhs.block || index != rhs.index;
     }
   };
 
   static void destroy(BlockStore<T, BLOCK_SIZE, REVERSE_ORDER> *block_store);
 
   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 {
         new_block->next = nullptr;
         current->next = new_block;
       }
       current = new_block;
       fill_count = 0;
     }
     T *obj = reinterpret_cast<T *>(current->data + fill_count * sizeof(T));
     ++fill_count;
     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() {
     return *reinterpret_cast<T *>(getLastBlock()->data +
                                   sizeof(T) * (fill_count - 1));
   }
 
   void pop_back() {
     fill_count--;
     if (fill_count || current == &first)
       return;
     auto [last, prev] = getLastBlocks();
     if (REVERSE_ORDER) {
       assert(last == current);
       current = current->next;
     } else {
       assert(prev->next == last);
       current = prev;
       current->next = nullptr;
     }
     if (last != &first)
-      ::free(last);
+      delete last;
     fill_count = BLOCK_SIZE;
   }
 
   bool empty() const { return current == &first && !fill_count; }
 
   iterator begin() {
     if (REVERSE_ORDER)
       return iterator(current, fill_count);

 template <typename T, size_t BLOCK_SIZE, bool REVERSE_ORDER>
 void BlockStore<T, BLOCK_SIZE, REVERSE_ORDER>::destroy(
     BlockStore<T, BLOCK_SIZE, REVERSE_ORDER> *block_store) {
   if (REVERSE_ORDER) {
     auto current = block_store->current;
     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;
   block_store->fill_count = 0;
 }
 
 // A convenience type for reverse order block stores.
 template <typename T, size_t BLOCK_SIZE>
 using ReverseOrderBlockStore = BlockStore<T, BLOCK_SIZE, true>;
 
 } // namespace cpp
 } // namespace __llvm_libc
 
 #endif // LLVM_LIBC_SUPPORT_BLOCKSTORE_H

libc/src/stdlib/atexit.cpp

while (!exit_callbacks.empty()) {
     handler_list_mtx.lock();
   }
   ExitCallbackList::destroy(&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;
 }
 
 // TODO: Handle the last dso handle argument.
 extern "C" int __cxa_atexit(AtExitCallback *callback, void *payload, void *) {
   return add_atexit_unit({callback, payload});
 }
 
 LLVM_LIBC_FUNCTION(int, atexit, (StdCAtExitCallback * callback)) {
   return add_atexit_unit(
       {&stdc_at_exit_func, reinterpret_cast<void *>(callback)});
 }
 
 } // namespace __llvm_libc

libc/test/src/__support/blockstore_test.cpp

 };
 
 class LlvmLibcBlockStoreTest : public __llvm_libc::testing::Test {
 public:
   template <size_t BLOCK_SIZE, size_t ELEMENT_COUNT, bool REVERSE>
   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) {
       Element &e = *iter;
       if (REVERSE) {
         int j = ELEMENT_COUNT - 1 - i;
         ASSERT_EQ(e.a, j);
         ASSERT_EQ(e.b, long(j * 2));
         ASSERT_EQ(e.c, unsigned(j * 3));
       } else {
         ASSERT_EQ(e.a, i);
         ASSERT_EQ(e.b, long(i * 2));
         ASSERT_EQ(e.c, unsigned(i * 3));
       }
     }
     ASSERT_EQ(i, int(ELEMENT_COUNT));
     __llvm_libc::cpp::BlockStore<Element, BLOCK_SIZE, REVERSE>::destroy(
         &block_store);
   }
 
   template <bool REVERSE> void back_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);
   }
 
   template <bool REVERSE> void empty_test() {
     using __llvm_libc::cpp::BlockStore;
     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);
   }
 };
 
 TEST_F(LlvmLibcBlockStoreTest, PopulateAndIterate4) {
   populate_and_iterate<4, 4, false>();
 }