From 5e3674ca805202eee22829ee936c11a38dabd6ab Mon Sep 17 00:00:00 2001
From: Jonathan Coens <jcoens@fb.com>
Date: Thu, 4 Apr 2013 15:25:52 -0700
Subject: [PATCH] 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
---
 folly/Arena-inl.h        |  7 +++++++
 folly/Arena.h            | 13 +++++++++----
 folly/test/ArenaTest.cpp | 11 +++++++++++
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/folly/Arena-inl.h b/folly/Arena-inl.h
index 719a35ec..d97ea697 100644
--- 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 bb43f45e..a8d8d7f8 100644
--- 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 493cd148..9fdc6f76 100644
--- 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);
-- 
2.34.1