Add clear() API to folly::Arena

Summary: Add a `clear()` API to render the Arena like new, except that previously-allocated blocks are preserved to be reused in the future.  (Caveat: we do not reuse "large blocks")

Reviewed By: yfeldblum, davidtgoldblatt

Differential Revision: D22209898

fbshipit-source-id: c2f27d70ee09855df8c8c6fc6d162fca92b9eac8

folly/memory/Arena-inl.h

@@ -42,7 +42,7 @@ void* Arena<Alloc>::allocateSlow(size_t size) {
     void* mem = AllocTraits::allocate(alloc(), allocSize);
     auto blk = new (mem) LargeBlock(allocSize);
     start = blk->start();
-    largeBlocks_.push_front(*blk);
+    largeBlocks_.push_back(*blk);
   } else {
     // Allocate a normal sized block and carve out size bytes from it
     // Will allocate more than size bytes if convenient
@@ -52,7 +52,8 @@ void* Arena<Alloc>::allocateSlow(size_t size) {
     void* mem = AllocTraits::allocate(alloc(), allocSize);
     auto blk = new (mem) Block();
     start = blk->start();
-    blocks_.push_front(*blk);
+    blocks_.push_back(*blk);
+    currentBlock_ = blocks_.last();
     ptr_ = start + size;
     end_ = start + allocSize - sizeof(Block);
     assert(ptr_ <= end_);
@@ -75,18 +76,4 @@ void Arena<Alloc>::merge(Arena<Alloc>&& other) {
   totalAllocatedSize_ += other.totalAllocatedSize_;
   other.totalAllocatedSize_ = 0;
 }
-
-template <class Alloc>
-Arena<Alloc>::~Arena() {
-  blocks_.clear_and_dispose([this](Block* b) {
-    b->~Block();
-    AllocTraits::deallocate(alloc(), b, blockGoodAllocSize());
-  });
-  largeBlocks_.clear_and_dispose([this](LargeBlock* b) {
-    auto size = b->allocSize;
-    b->~LargeBlock();
-    AllocTraits::deallocate(alloc(), b, size);
-  });
-}
-
 } // namespace folly

folly/memory/Arena.h

@@ -66,6 +66,7 @@ class Arena {
       size_t sizeLimit = kNoSizeLimit,
       size_t maxAlign = kDefaultMaxAlign)
       : allocAndSize_(alloc, minBlockSize),
+        currentBlock_(blocks_.last()),
         ptr_(nullptr),
         end_(nullptr),
         totalAllocatedSize_(0),
@@ -78,7 +79,10 @@ class Arena {
     }
   }
 
-  ~Arena();
+  ~Arena() {
+    freeBlocks();
+    freeLargeBlocks();
+  }
 
   void* allocate(size_t size) {
     size = roundUp(size);
@@ -93,6 +97,16 @@ class Arena {
       return r;
     }
 
+    if (canReuseExistingBlock(size)) {
+      currentBlock_++;
+      char* r = currentBlock_->start();
+      ptr_ = r + size;
+      end_ = r + blockGoodAllocSize() - sizeof(Block);
+      assert(ptr_ <= end_);
+      assert(isAligned(r));
+      return r;
+    }
+
     // Not enough room in the current block
     void* r = allocateSlow(size);
     assert(isAligned(r));
@@ -106,6 +120,19 @@ class Arena {
   // Transfer ownership of all memory allocated from "other" to "this".
   void merge(Arena&& other);
 
+  void clear() {
+    freeLargeBlocks(); // We don't reuse large blocks
+    if (blocks_.empty()) {
+      return;
+    }
+    bytesUsed_ = 0;
+    currentBlock_ = blocks_.begin();
+    char* start = currentBlock_->start();
+    ptr_ = start;
+    end_ = start + blockGoodAllocSize() - sizeof(Block);
+    assert(ptr_ <= end_);
+  }
+
   // Gets the total memory used by the arena
   size_t totalSize() const {
     return totalAllocatedSize_ + sizeof(Arena);
@@ -157,6 +184,33 @@ class Arena {
     ~LargeBlock() = default;
   };
 
+  bool canReuseExistingBlock(size_t size) {
+    if (size > minBlockSize()) {
+      // We don't reuse large blocks
+      return false;
+    }
+    if (blocks_.empty() || currentBlock_ == blocks_.last()) {
+      // No regular blocks to reuse
+      return false;
+    }
+    return true;
+  }
+
+  void freeBlocks() {
+    blocks_.clear_and_dispose([this](Block* b) {
+      b->~Block();
+      AllocTraits::deallocate(alloc(), b, blockGoodAllocSize());
+    });
+  }
+
+  void freeLargeBlocks() {
+    largeBlocks_.clear_and_dispose([this](LargeBlock* b) {
+      auto size = b->allocSize;
+      b->~LargeBlock();
+      AllocTraits::deallocate(alloc(), b, size);
+    });
+  }
+
  public:
   static constexpr size_t kDefaultMinBlockSize = 4096 - sizeof(Block);
   static constexpr size_t kNoSizeLimit = 0;
@@ -220,6 +274,7 @@ class Arena {
 
   AllocAndSize allocAndSize_;
   BlockList blocks_;
+  typename BlockList::iterator currentBlock_;
   LargeBlockList largeBlocks_;
   char* ptr_;
   char* end_;

folly/memory/test/ArenaTest.cpp

@@ -186,6 +186,42 @@ TEST(Arena, ExtremeSize) {
   EXPECT_THROW(arena.allocate(SIZE_MAX - 2), std::bad_alloc);
 }
 
+TEST(Arena, Clear) {
+  static const size_t blockSize = 1024;
+  SysArena arena(blockSize);
+
+  for (int i = 0; i < 10; ++i) {
+    std::vector<size_t> sizes(1000);
+    std::generate(sizes.begin(), sizes.end(), []() {
+      return std::rand() % blockSize * 2;
+    });
+
+    std::vector<void*> addresses;
+    for (auto s : sizes) {
+      addresses.push_back(arena.allocate(s));
+    }
+
+    const size_t totalSize = arena.totalSize();
+    const size_t bytesUsed = arena.bytesUsed();
+
+    arena.clear();
+
+    int j = 0;
+    for (auto s : sizes) {
+      auto addr = arena.allocate(s);
+      if (s <= blockSize) {
+        EXPECT_EQ(addr, addresses[j]);
+      }
+      j++;
+    }
+
+    EXPECT_EQ(arena.bytesUsed(), bytesUsed);
+    EXPECT_GT(arena.totalSize(), totalSize);
+
+    arena.clear();
+  }
+}
+
 int main(int argc, char* argv[]) {
   testing::InitGoogleTest(&argc, argv);
   gflags::ParseCommandLineFlags(&argc, &argv, true);