From b197e2c78efba32bb05823158b85815b7f5966d8 Mon Sep 17 00:00:00 2001
From: Roman Gershman <romange@gmail.com>
Date: Fri, 11 Mar 2022 09:59:35 +0200
Subject: [PATCH] Throw bad_alloc exception in MiMemoryResource if the
 allocation fails

---
 src/core/dash.h               |  1 +
 src/core/dash_test.cc         | 46 +++++++++++++++++++++++++++++++++++
 src/core/mi_memory_resource.h |  6 +++--
 3 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/src/core/dash.h b/src/core/dash.h
index 66359ea..dc4d606 100644
--- a/src/core/dash.h
+++ b/src/core/dash.h
@@ -339,6 +339,7 @@ DashTable<_Key, _Value, Policy>::DashTable(size_t capacity_log, const Policy& po
   segment_.resize(unique_segments_);
   std::pmr::polymorphic_allocator<SegmentType> pa(mr);
 
+  // I assume we have enough memory to create the initial table and do not check allocations.
   for (auto& ptr : segment_) {
     ptr = pa.allocate(1);
     pa.construct(ptr, global_depth_);  //   new SegmentType(global_depth_);
diff --git a/src/core/dash_test.cc b/src/core/dash_test.cc
index 29a57b7..7e5d47d 100644
--- a/src/core/dash_test.cc
+++ b/src/core/dash_test.cc
@@ -71,6 +71,39 @@ struct UInt64Policy : public BasicDashPolicy {
   }
 };
 
+class CappedResource final : public std::pmr::memory_resource {
+ public:
+  explicit CappedResource(size_t cap) : cap_(cap) {
+  }
+
+  size_t used() const {
+    return used_;
+  }
+
+ private:
+  void* do_allocate(std::size_t size, std::size_t align) {
+    if (used_ + size > cap_)
+      throw std::bad_alloc{};
+
+    void* res = pmr::get_default_resource()->allocate(size, align);
+    used_ += size;
+
+    return res;
+  }
+
+  void do_deallocate(void* ptr, std::size_t size, std::size_t align) {
+    used_ -= size;
+    pmr::get_default_resource()->deallocate(ptr, size, align);
+  }
+
+  bool do_is_equal(const std::pmr::memory_resource& o) const noexcept {
+    return this == &o;
+  }
+
+  size_t cap_;
+  size_t used_ = 0;
+};
+
 using Segment = detail::Segment<uint64_t, Buf24>;
 using Dash64 = DashTable<uint64_t, uint64_t, UInt64Policy>;
 
@@ -296,6 +329,19 @@ TEST_F(DashTest, Insert2) {
   }
 }
 
+TEST_F(DashTest, InsertOOM) {
+  CappedResource resource(1 << 15);
+  Dash64 dt{1, UInt64Policy{}, &resource};
+
+  ASSERT_THROW(
+      {
+        for (size_t i = 0; i < (1 << 14); ++i) {
+          dt.Insert(i, 0);
+        }
+      },
+      std::bad_alloc);
+}
+
 struct Item {
   char buf[24];
 };
diff --git a/src/core/mi_memory_resource.h b/src/core/mi_memory_resource.h
index 572bfd2..3a584e4 100644
--- a/src/core/mi_memory_resource.h
+++ b/src/core/mi_memory_resource.h
@@ -22,8 +22,10 @@ class MiMemoryResource final : public std::pmr::memory_resource {
  private:
   void* do_allocate(std::size_t size, std::size_t align) {
     void* res = mi_heap_malloc_aligned(heap_, size, align);
-    if (res)
-      used_ += mi_good_size(size);
+
+    if (!res)
+      throw std::bad_alloc{};
+    used_ += mi_good_size(size);
     return res;
   }