Harden failure signal handler in the face of memory corruptions

[folly.git] / folly / Arena.h

diff --git a/folly/Arena.h b/folly/Arena.h
index 8db3ca5dd0b19f28e05bda4868230219e8e9bdbe..c119d8c16ee762f916d913cea4c6bf639a42b697 100644 (file)
--- a/folly/Arena.h
+++ b/folly/Arena.h
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012 Facebook, Inc.
+ * Copyright 2014 Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -60,16 +60,21 @@ template <class Alloc>
 class Arena {
  public:
   explicit Arena(const Alloc& alloc,
-                 size_t minBlockSize = kDefaultMinBlockSize)
-    : allocAndSize_(alloc, minBlockSize),
-      ptr_(nullptr),
-      end_(nullptr) {
+                 size_t minBlockSize = kDefaultMinBlockSize,
+                 size_t sizeLimit = 0)
+    : allocAndSize_(alloc, minBlockSize)
+    , ptr_(nullptr)
+    , end_(nullptr)
+    , totalAllocatedSize_(0)
+    , bytesUsed_(0)
+    , sizeLimit_(sizeLimit) {
   }
 
   ~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
@@ -92,6 +97,19 @@ class Arena {
   // Transfer ownership of all memory allocated from "other" to "this".
   void merge(Arena&& other);
 
+  // Gets the total memory used by the arena
+  size_t totalSize() const {
+    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;
@@ -174,6 +192,9 @@ class Arena {
   BlockList blocks_;
   char* ptr_;
   char* end_;
+  size_t totalAllocatedSize_;
+  size_t bytesUsed_;
+  size_t sizeLimit_;
 };
 
 /**
@@ -193,9 +214,7 @@ struct ArenaAllocatorTraits {
 class SysAlloc {
  public:
   void* allocate(size_t size) {
-    void* mem = malloc(size);
-    if (!mem) throw std::bad_alloc();
-    return mem;
+    return checkedMalloc(size);
   }
 
   void deallocate(void* p) {
@@ -215,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) {
   }
 };