Wrap all stl containers with a marl::StlAllocator

`StlAllocator` uses a `marl::Allocator` for performing the allocations,
avoiding untracked heap allocations.

Issue: google#131

include/marl/containers.h

@@ -22,18 +22,66 @@
 #include <cstddef>    // size_t
 #include <utility>    // std::move
 
+#include <deque>
+#include <map>
+#include <set>
+#include <unordered_map>
+#include <unordered_set>
+
 namespace marl {
 namespace containers {
 
+////////////////////////////////////////////////////////////////////////////////
+// STL wrappers
+// STL containers that use a marl::StlAllocator backed by a marl::Allocator.
+// Note: These may be re-implemented to optimize for marl's usage cases.
+// See: https://github.com/google/marl/issues/129
+////////////////////////////////////////////////////////////////////////////////
+template <typename T>
+using deque = std::deque<T, StlAllocator<T>>;
+
+template <typename K, typename V, typename C = std::less<K>>
+using map = std::map<K, V, C, StlAllocator<std::pair<const K, V>>>;
+
+template <typename K, typename C = std::less<K>>
+using set = std::set<K, C, StlAllocator<K>>;
+
+template <typename K,
+          typename V,
+          typename H = std::hash<K>,
+          typename E = std::equal_to<K>>
+using unordered_map =
+    std::unordered_map<K, V, H, E, StlAllocator<std::pair<const K, V>>>;
+
+template <typename K, typename H = std::hash<K>, typename E = std::equal_to<K>>
+using unordered_set = std::unordered_set<K, H, E, StlAllocator<K>>;
+
+// take() takes and returns the front value from the deque.
+template <typename T>
+inline T take(deque<T>& queue) {
+  auto out = std::move(queue.front());
+  queue.pop_front();
+  return out;
+}
+
+// take() takes and returns the first value from the unordered_set.
+template <typename T, typename H, typename E>
+inline T take(unordered_set<T, H, E>& set) {
+  auto it = set.begin();
+  auto out = std::move(*it);
+  set.erase(it);
+  return out;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // vector<T, BASE_CAPACITY>
 ////////////////////////////////////////////////////////////////////////////////
 
 // vector is a container of contiguously stored elements.
-// Unlike std::vector, marl::containers::vector keeps the first BASE_CAPACITY
-// elements internally, which will avoid dynamic heap allocations.
-// Once the vector exceeds BASE_CAPACITY elements, vector will allocate storage
-// from the heap.
+// Unlike std::vector, marl::containers::vector keeps the first
+// BASE_CAPACITY elements internally, which will avoid dynamic heap
+// allocations. Once the vector exceeds BASE_CAPACITY elements, vector will
+// allocate storage from the heap.
 template <typename T, int BASE_CAPACITY>
 class vector {
  public:
@@ -74,6 +122,10 @@ class vector {
   inline size_t cap() const;
   inline void resize(size_t n);
   inline void reserve(size_t n);
+  inline T* data();
+  inline const T* data() const;
+
+  Allocator* const allocator;
 
  private:
   using TStorage = typename marl::aligned_storage<sizeof(T), alignof(T)>::type;
@@ -82,7 +134,6 @@ class vector {
 
   inline void free();
 
-  Allocator* const allocator;
   size_t count = 0;
   size_t capacity = BASE_CAPACITY;
   TStorage buffer[BASE_CAPACITY];
@@ -272,6 +323,16 @@ void vector<T, BASE_CAPACITY>::reserve(size_t n) {
   }
 }
 
+template <typename T, int BASE_CAPACITY>
+T* vector<T, BASE_CAPACITY>::data() {
+  return elements;
+}
+
+template <typename T, int BASE_CAPACITY>
+const T* vector<T, BASE_CAPACITY>::data() const {
+  return elements;
+}
+
 template <typename T, int BASE_CAPACITY>
 void vector<T, BASE_CAPACITY>::free() {
   for (size_t i = 0; i < count; i++) {

include/marl/memory.h

@@ -27,6 +27,9 @@
 
 namespace marl {
 
+template <typename T>
+struct StlAllocator;
+
 // pageSize() returns the size in bytes of a virtual memory page for the host
 // system.
 size_t pageSize();
@@ -36,6 +39,19 @@ inline T alignUp(T val, T alignment) {
   return alignment * ((val + alignment - 1) / alignment);
 }
 
+// aligned_storage() is a replacement for std::aligned_storage that isn't busted
+// on older versions of MSVC.
+template <size_t SIZE, size_t ALIGNMENT>
+struct aligned_storage {
+  struct alignas(ALIGNMENT) type {
+    unsigned char data[SIZE];
+  };
+};
+
+///////////////////////////////////////////////////////////////////////////////
+// Allocation
+///////////////////////////////////////////////////////////////////////////////
+
 // Allocation holds the result of a memory allocation from an Allocator.
 struct Allocation {
   // Intended usage of the allocation. Used for allocation trackers.
@@ -45,6 +61,7 @@ struct Allocation {
     Create,  // Allocator::create(), make_unique(), make_shared()
     Vector,  // marl::containers::vector<T>
     List,    // marl::containers::list<T>
+    Stl,     // marl::StlAllocator
     Count,   // Not intended to be used as a usage type - used for upper bound.
   };
 
@@ -60,6 +77,10 @@ struct Allocation {
   Request request;      // Request used for the allocation.
 };
 
+///////////////////////////////////////////////////////////////////////////////
+// Allocator
+///////////////////////////////////////////////////////////////////////////////
+
 // Allocator is an interface to a memory allocator.
 // Marl provides a default implementation with Allocator::Default.
 class Allocator {
@@ -183,14 +204,9 @@ std::shared_ptr<T> Allocator::make_shared(ARGS&&... args) {
   return std::shared_ptr<T>(reinterpret_cast<T*>(alloc.ptr), Deleter{this});
 }
 
-// aligned_storage() is a replacement for std::aligned_storage that isn't busted
-// on older versions of MSVC.
-template <size_t SIZE, size_t ALIGNMENT>
-struct aligned_storage {
-  struct alignas(ALIGNMENT) type {
-    unsigned char data[SIZE];
-  };
-};
+///////////////////////////////////////////////////////////////////////////////
+// TrackedAllocator
+///////////////////////////////////////////////////////////////////////////////
 
 // TrackedAllocator wraps an Allocator to track the allocations made.
 class TrackedAllocator : public Allocator {
@@ -280,6 +296,103 @@ void TrackedAllocator::free(const Allocation& allocation) {
   return allocator->free(allocation);
 }
 
+///////////////////////////////////////////////////////////////////////////////
+// StlAllocator
+///////////////////////////////////////////////////////////////////////////////
+
+// StlAllocator exposes an STL-compatible allocator for a given a
+// marl::Allocator.
+template <typename T>
+struct StlAllocator {
+  using value_type = T;
+  using pointer = T*;
+  using const_pointer = const T*;
+  using reference = T&;
+  using const_reference = const T&;
+  using size_type = size_t;
+  using difference_type = size_t;
+
+  // An equivalent STL allocator for a different type.
+  template <class U>
+  struct rebind {
+    typedef StlAllocator<U> other;
+  };
+
+  // Creation of the allocator is allowed, but if anyone were to try to use
+  // an object with a null allocator, it would fail. This however allows us to
+  // default-construct a bunch of objects in a container, and fill, in their
+  // allocators on first use.
+  // StlAllocator() : allocator(nullptr) {}
+
+  // All allocations will be done through this allocator. It must remain
+  // valid until this StlAllocator and all allocators created
+  // from it have been destroyed.
+  StlAllocator(Allocator* allocator) : allocator(allocator) {}
+
+  template <typename U>
+  StlAllocator(const StlAllocator<U>& other) {
+    allocator = other.allocator;
+  }
+
+  // Returns the actual address of x even in presence of overloaded operator&.
+  pointer address(reference x) const { return &x; }
+  const_pointer address(const_reference x) const { return &x; }
+
+  // Allocates the memory for n objects of type T. Does not
+  // actually construct the objects.
+  T* allocate(std::size_t n) {
+    auto alloc = allocator->allocate(request<T>(n));
+    return reinterpret_cast<T*>(alloc.ptr);
+  }
+
+  // Deallocates the memory for n Objects of size T.
+  void deallocate(T* p, std::size_t n) {
+    Allocation alloc;
+    alloc.ptr = p;
+    alloc.request = request<T>(n);
+    allocator->free(alloc);
+  }
+
+  // Returns the maximum theoretically possible number of T stored in this
+  // allocator.
+  size_type max_size() const {
+    return std::numeric_limits<size_type>::max() / sizeof(value_type);
+  }
+
+  // Copy constructs an object of type T at the location given by p.
+  void construct(pointer p, const_reference val) { new (p) T(val); }
+
+  // Constructs an object of Type U at the location given by P passing
+  // through all other arguments to the constructor.
+  template <typename U, typename... Args>
+  void construct(U* p, Args&&... args) {
+    ::new ((void*)p) U(std::forward<Args>(args)...);
+  }
+
+  // Deconstructs the object at p. It does not free the memory.
+  void destroy(pointer p) { ((T*)p)->~T(); }
+
+  // Deconstructs the object at p. It does not free the memory.
+  template <typename U>
+  void destroy(U* p) {
+    p->~U();
+  }
+
+ private:
+  template <typename T>
+  Allocation::Request request(size_t n) const {
+    Allocation::Request req = {};
+    req.size = sizeof(T) * n;
+    req.alignment = alignof(T);
+    req.usage = Allocation::Usage::Stl;
+    return req;
+  }
+
+  template <typename U>
+  friend struct StlAllocator;
+  Allocator* allocator;
+};
+
 }  // namespace marl
 
 #endif  // marl_memory_h

include/marl/pool.h

@@ -363,7 +363,7 @@ class UnboundedPool : public Pool<T> {
 
     Allocator* allocator;
     marl::mutex mutex;
-    std::vector<Item*> items;
+    containers::vector<Item*, 4> items;
     Item* free = nullptr;
   };
 
@@ -373,7 +373,7 @@ class UnboundedPool : public Pool<T> {
 
 template <typename T, PoolPolicy POLICY>
 UnboundedPool<T, POLICY>::Storage::Storage(Allocator* allocator)
-    : allocator(allocator) {}
+    : allocator(allocator), items(allocator) {}
 
 template <typename T, PoolPolicy POLICY>
 UnboundedPool<T, POLICY>::Storage::~Storage() {

include/marl/scheduler.h

@@ -15,6 +15,7 @@
 #ifndef marl_scheduler_h
 #define marl_scheduler_h
 
+#include "containers.h"
 #include "debug.h"
 #include "deprecated.h"
 #include "memory.h"
@@ -26,14 +27,8 @@
 #include <atomic>
 #include <chrono>
 #include <condition_variable>
-#include <deque>
 #include <functional>
-#include <map>
-#include <set>
 #include <thread>
-#include <unordered_map>
-#include <unordered_set>
-#include <vector>
 
 namespace marl {
 
@@ -291,6 +286,8 @@ class Scheduler {
 
   // WaitingFibers holds all the fibers waiting on a timeout.
   struct WaitingFibers {
+    inline WaitingFibers(Allocator*);
+
     // operator bool() returns true iff there are any wait fibers.
     inline operator bool() const;
 
@@ -317,15 +314,15 @@ class Scheduler {
       Fiber* fiber;
       inline bool operator<(const Timeout&) const;
     };
-    std::set<Timeout> timeouts;
-    std::unordered_map<Fiber*, TimePoint> fibers;
+    containers::set<Timeout, std::less<Timeout>> timeouts;
+    containers::unordered_map<Fiber*, TimePoint> fibers;
   };
 
   // TODO: Implement a queue that recycles elements to reduce number of
   // heap allocations.
-  using TaskQueue = std::deque<Task>;
-  using FiberQueue = std::deque<Fiber*>;
-  using FiberSet = std::unordered_set<Fiber*>;
+  using TaskQueue = containers::deque<Task>;
+  using FiberQueue = containers::deque<Fiber*>;
+  using FiberSet = containers::unordered_set<Fiber*>;
 
   // Workers executes Tasks on a single thread.
   // Once a task is started, it may yield to other tasks on the same Worker.
@@ -437,6 +434,8 @@ class Scheduler {
 
     // Work holds tasks and fibers that are enqueued on the Worker.
     struct Work {
+      inline Work(Allocator*);
+
       std::atomic<uint64_t> num = {0};  // tasks.size() + fibers.size()
       GUARDED_BY(mutex) uint64_t numBlockedFibers = 0;
       GUARDED_BY(mutex) TaskQueue tasks;
@@ -474,7 +473,7 @@ class Scheduler {
     Thread thread;
     Work work;
     FiberSet idleFibers;  // Fibers that have completed which can be reused.
-    std::vector<Allocator::unique_ptr<Fiber>>
+    containers::vector<Allocator::unique_ptr<Fiber>, 16>
         workerFibers;  // All fibers created by this worker.
     FastRnd rng;
     bool shutdown = false;
@@ -506,8 +505,11 @@ class Scheduler {
   std::array<Worker*, MaxWorkerThreads> workerThreads;
 
   struct SingleThreadedWorkers {
+    inline SingleThreadedWorkers(Allocator*);
+
     using WorkerByTid =
-        std::unordered_map<std::thread::id, Allocator::unique_ptr<Worker>>;
+        containers::unordered_map<std::thread::id,
+                                  Allocator::unique_ptr<Worker>>;
     marl::mutex mutex;
     GUARDED_BY(mutex) std::condition_variable unbind;
     GUARDED_BY(mutex) WorkerByTid byTid;