Add bytesUsed() to folly::Arena

Summary:
Returns number of bytes "really used by the user", i.e. number of bytes
allocated on the arena via calls to `allocate()`.

Test Plan: compiled, unit tests

Reviewed By: azzolini@fb.com

FB internal diff: D636976

diff --git a/folly/Arena.h b/folly/Arena.h
index 26917f0c16d1ca3202b427302bad7a083883ddef..3ddcea9f47b7ee2c3d7d241352e09c273458dbcd 100644 (file)
--- a/folly/Arena.h
+++ b/folly/Arena.h
@@ -61,16 +61,18 @@ class Arena {
  public:
   explicit Arena(const Alloc& alloc,
                  size_t minBlockSize = kDefaultMinBlockSize)
-    : allocAndSize_(alloc, minBlockSize),
-      ptr_(nullptr),
-      end_(nullptr),
-      totalAllocatedSize_(0) {
+    : allocAndSize_(alloc, minBlockSize)
+    , ptr_(nullptr)
+    , end_(nullptr)
+    , totalAllocatedSize_(0)
+    , bytesUsed_(0) {
   }
 
   ~Arena();
 
   void* allocate(size_t size) {
     size = roundUp(size);
+    bytesUsed_ += size;
 
     if (LIKELY(end_ - ptr_ >= size)) {
       // Fast path: there's enough room in the current block
@@ -98,6 +100,14 @@ class Arena {
     return totalAllocatedSize_ + sizeof(Arena);
   }
 
+  // Gets the total number of "used" bytes, i.e. bytes that the arena users
+  // allocated via the calls to `allocate`. Doesn't include fragmentation, e.g.
+  // if block size is 4KB and you allocate 2 objects of 3KB in size,
+  // `bytesUsed()` will be 6KB, while `totalSize()` will be 8KB+.
+  size_t bytesUsed() const {
+    return bytesUsed_;
+  }
+
  private:
   // not copyable
   Arena(const Arena&) = delete;
@@ -181,6 +191,7 @@ class Arena {
   char* ptr_;
   char* end_;
   size_t totalAllocatedSize_;
+  size_t bytesUsed_;
 };
 
 /**

diff --git a/folly/test/ArenaTest.cpp b/folly/test/ArenaTest.cpp
index cfbd7dcfd86a5c1afe77c2a088004befd5c7b916..5c59bc0479a8767072dfcc2ff6225c82e0289ffa 100644 (file)
--- a/folly/test/ArenaTest.cpp
+++ b/folly/test/ArenaTest.cpp
@@ -84,6 +84,46 @@ TEST(Arena, SizeSanity) {
           << maximum_size;
 }
 
+TEST(Arena, BytesUsedSanity) {
+  static const size_t smallChunkSize = 1024;
+  static const size_t blockSize = goodMallocSize(16 * smallChunkSize);
+  const size_t bigChunkSize = blockSize - 4 * smallChunkSize;
+
+  size_t bytesUsed = 0;
+
+  SysArena arena(blockSize);
+  EXPECT_EQ(arena.bytesUsed(), bytesUsed);
+
+  // Insert 2 small chunks
+  arena.allocate(smallChunkSize);
+  arena.allocate(smallChunkSize);
+  bytesUsed += 2 * smallChunkSize;
+  EXPECT_EQ(arena.bytesUsed(), bytesUsed);
+  EXPECT_TRUE(arena.totalSize() >= blockSize);
+  EXPECT_TRUE(arena.totalSize() <= 2 * blockSize);
+
+  // Insert big chunk, should still fit in one block
+  arena.allocate(bigChunkSize);
+  bytesUsed += bigChunkSize;
+  EXPECT_EQ(arena.bytesUsed(), bytesUsed);
+  EXPECT_TRUE(arena.totalSize() >= blockSize);
+  EXPECT_TRUE(arena.totalSize() <= 2 * blockSize);
+
+  // Insert big chunk once more, should trigger new block allocation
+  arena.allocate(bigChunkSize);
+  bytesUsed += bigChunkSize;
+  EXPECT_EQ(arena.bytesUsed(), bytesUsed);
+  EXPECT_TRUE(arena.totalSize() >= 2 * blockSize);
+  EXPECT_TRUE(arena.totalSize() <= 3 * blockSize);
+
+  // Test that bytesUsed() accounts for alignment
+  static const size_t tinyChunkSize = 7;
+  arena.allocate(tinyChunkSize);
+  EXPECT_TRUE(arena.bytesUsed() >= bytesUsed + tinyChunkSize);
+  size_t delta = arena.bytesUsed() - bytesUsed;
+  EXPECT_EQ(delta & (delta - 1), 0);
+}
+
 TEST(Arena, Vector) {
   static const size_t requestedBlockSize = 64;
   SysArena arena(requestedBlockSize);