Soft-limit for arenas

Summary: Create an artificial limit on an arena to start throwing bad_alloc before running out of system memory

Test Plan: adjust unit test

Reviewed By: marcelo.juchem@fb.com

FB internal diff: D762695

diff --git a/folly/Arena-inl.h b/folly/Arena-inl.h
index 719a35ec9e76302565549b42dd801db605fbb39f..d97ea6973700f42431d4ab76975c67d220f4d8b4 100644 (file)
--- a/folly/Arena-inl.h
+++ b/folly/Arena-inl.h
@@ -45,6 +45,13 @@ template <class Alloc>
 void* Arena<Alloc>::allocateSlow(size_t size) {
   std::pair<Block*, size_t> p;
   char* start;
+
+
+  size_t allocSize = std::max(size, minBlockSize()) + sizeof(Block);
+  if(sizeLimit_ && allocSize > sizeLimit_ - totalAllocatedSize_) {
+    throw std::bad_alloc();
+  }
+
   if (size > minBlockSize()) {
     // Allocate a large block for this chunk only, put it at the back of the
     // list so it doesn't get used for small allocations; don't change ptr_

diff --git a/folly/Arena.h b/folly/Arena.h
index bb43f45e74bc3c4a294daf84b5c99c108982336e..a8d8d7f806a2a1dae5e19e05cf1facd46940efee 100644 (file)
--- a/folly/Arena.h
+++ b/folly/Arena.h
@@ -60,12 +60,14 @@ template <class Alloc>
 class Arena {
  public:
   explicit Arena(const Alloc& alloc,
-                 size_t minBlockSize = kDefaultMinBlockSize)
+                 size_t minBlockSize = kDefaultMinBlockSize,
+                 size_t sizeLimit = 0)
     : allocAndSize_(alloc, minBlockSize)
     , ptr_(nullptr)
     , end_(nullptr)
     , totalAllocatedSize_(0)
-    , bytesUsed_(0) {
+    , bytesUsed_(0)
+    , sizeLimit_(sizeLimit) {
   }
 
   ~Arena();
@@ -192,6 +194,7 @@ class Arena {
   char* end_;
   size_t totalAllocatedSize_;
   size_t bytesUsed_;
+  size_t sizeLimit_;
 };
 
 /**
@@ -231,8 +234,10 @@ struct ArenaAllocatorTraits<SysAlloc> {
  */
 class SysArena : public Arena<SysAlloc> {
  public:
-  explicit SysArena(size_t minBlockSize = kDefaultMinBlockSize)
-    : Arena<SysAlloc>(SysAlloc(), minBlockSize) {
+  explicit SysArena(
+    size_t minBlockSize = kDefaultMinBlockSize,
+    size_t sizeLimit = 0)
+      : Arena<SysAlloc>(SysAlloc(), minBlockSize, sizeLimit) {
   }
 };
 

diff --git a/folly/test/ArenaTest.cpp b/folly/test/ArenaTest.cpp
index 493cd148b9cf1d6a5b2f7917ec41149c2372d558..9fdc6f766d2897cee0d1d817398d80470d191494 100644 (file)
--- a/folly/test/ArenaTest.cpp
+++ b/folly/test/ArenaTest.cpp
@@ -142,6 +142,17 @@ TEST(Arena, Vector) {
   }
 }
 
+TEST(Arena, SizeLimit) {
+  static const size_t requestedBlockSize = sizeof(size_t);
+  static const size_t maxSize = 10 * requestedBlockSize;
+
+  SysArena arena(requestedBlockSize, maxSize);
+
+  void* a = arena.allocate(sizeof(size_t));
+  EXPECT_TRUE(a != nullptr);
+  EXPECT_THROW(arena.allocate(maxSize + 1), std::bad_alloc);
+}
+
 int main(int argc, char *argv[]) {
   testing::InitGoogleTest(&argc, argv);
   google::ParseCommandLineFlags(&argc, &argv, true);