1 // Copyright 2019 The Marl Authors.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #ifndef marl_memory_h
16 #define marl_memory_h
17 
18 #include "debug.h"
19 #include "export.h"
20 
21 #include <stdint.h>
22 
23 #include <array>
24 #include <cstdlib>
25 #include <memory>
26 #include <mutex>
27 #include <utility>  // std::forward
28 
29 namespace marl {
30 
31 template <typename T>
32 struct StlAllocator;
33 
34 // pageSize() returns the size in bytes of a virtual memory page for the host
35 // system.
36 MARL_EXPORT
37 size_t pageSize();
38 
39 template <typename T>
alignUp(T val,T alignment)40 MARL_NO_EXPORT inline T alignUp(T val, T alignment) {
41   return alignment * ((val + alignment - 1) / alignment);
42 }
43 
44 // aligned_storage() is a replacement for std::aligned_storage that isn't busted
45 // on older versions of MSVC.
46 template <size_t SIZE, size_t ALIGNMENT>
47 struct aligned_storage {
48   struct alignas(ALIGNMENT) type {
49     unsigned char data[SIZE];
50   };
51 };
52 
53 ///////////////////////////////////////////////////////////////////////////////
54 // Allocation
55 ///////////////////////////////////////////////////////////////////////////////
56 
57 // Allocation holds the result of a memory allocation from an Allocator.
58 struct Allocation {
59   // Intended usage of the allocation. Used for allocation trackers.
60   enum class Usage : uint8_t {
61     Undefined = 0,
62     Stack,   // Fiber stack
63     Create,  // Allocator::create(), make_unique(), make_shared()
64     Vector,  // marl::containers::vector<T>
65     List,    // marl::containers::list<T>
66     Stl,     // marl::StlAllocator
67     Count,   // Not intended to be used as a usage type - used for upper bound.
68   };
69 
70   // Request holds all the information required to make an allocation.
71   struct Request {
72     size_t size = 0;                 // The size of the allocation in bytes.
73     size_t alignment = 0;            // The minimum alignment of the allocation.
74     bool useGuards = false;          // Whether the allocation is guarded.
75     Usage usage = Usage::Undefined;  // Intended usage of the allocation.
76   };
77 
78   void* ptr = nullptr;  // The pointer to the allocated memory.
79   Request request;      // Request used for the allocation.
80 };
81 
82 ///////////////////////////////////////////////////////////////////////////////
83 // Allocator
84 ///////////////////////////////////////////////////////////////////////////////
85 
86 // Allocator is an interface to a memory allocator.
87 // Marl provides a default implementation with Allocator::Default.
88 class Allocator {
89  public:
90   // The default allocator. Initialized with an implementation that allocates
91   // from the OS. Can be assigned a custom implementation.
92   MARL_EXPORT static Allocator* Default;
93 
94   // Deleter is a smart-pointer compatible deleter that can be used to delete
95   // objects created by Allocator::create(). Deleter is used by the smart
96   // pointers returned by make_shared() and make_unique().
97   struct MARL_EXPORT Deleter {
98     MARL_NO_EXPORT inline Deleter();
99     MARL_NO_EXPORT inline Deleter(Allocator* allocator, size_t count);
100 
101     template <typename T>
102     MARL_NO_EXPORT inline void operator()(T* object);
103 
104     Allocator* allocator = nullptr;
105     size_t count = 0;
106   };
107 
108   // unique_ptr<T> is an alias to std::unique_ptr<T, Deleter>.
109   template <typename T>
110   using unique_ptr = std::unique_ptr<T, Deleter>;
111 
112   virtual ~Allocator() = default;
113 
114   // allocate() allocates memory from the allocator.
115   // The returned Allocation::request field must be equal to the Request
116   // parameter.
117   virtual Allocation allocate(const Allocation::Request&) = 0;
118 
119   // free() frees the memory returned by allocate().
120   // The Allocation must have all fields equal to those returned by allocate().
121   virtual void free(const Allocation&) = 0;
122 
123   // create() allocates and constructs an object of type T, respecting the
124   // alignment of the type.
125   // The pointer returned by create() must be deleted with destroy().
126   template <typename T, typename... ARGS>
127   inline T* create(ARGS&&... args);
128 
129   // destroy() destructs and frees the object allocated with create().
130   template <typename T>
131   inline void destroy(T* object);
132 
133   // make_unique() returns a new object allocated from the allocator wrapped
134   // in a unique_ptr that respects the alignemnt of the type.
135   template <typename T, typename... ARGS>
136   inline unique_ptr<T> make_unique(ARGS&&... args);
137 
138   // make_unique_n() returns an array of n new objects allocated from the
139   // allocator wrapped in a unique_ptr that respects the alignemnt of the
140   // type.
141   template <typename T, typename... ARGS>
142   inline unique_ptr<T> make_unique_n(size_t n, ARGS&&... args);
143 
144   // make_shared() returns a new object allocated from the allocator
145   // wrapped in a std::shared_ptr that respects the alignemnt of the type.
146   template <typename T, typename... ARGS>
147   inline std::shared_ptr<T> make_shared(ARGS&&... args);
148 
149  protected:
150   Allocator() = default;
151 };
152 
153 ///////////////////////////////////////////////////////////////////////////////
154 // Allocator::Deleter
155 ///////////////////////////////////////////////////////////////////////////////
Deleter()156 Allocator::Deleter::Deleter() : allocator(nullptr) {}
Deleter(Allocator * allocator,size_t count)157 Allocator::Deleter::Deleter(Allocator* allocator, size_t count)
158     : allocator(allocator), count(count) {}
159 
160 template <typename T>
operator()161 void Allocator::Deleter::operator()(T* object) {
162   object->~T();
163 
164   Allocation allocation;
165   allocation.ptr = object;
166   allocation.request.size = sizeof(T) * count;
167   allocation.request.alignment = alignof(T);
168   allocation.request.usage = Allocation::Usage::Create;
169   allocator->free(allocation);
170 }
171 
172 ///////////////////////////////////////////////////////////////////////////////
173 // Allocator
174 ///////////////////////////////////////////////////////////////////////////////
175 template <typename T, typename... ARGS>
create(ARGS &&...args)176 T* Allocator::create(ARGS&&... args) {
177   Allocation::Request request;
178   request.size = sizeof(T);
179   request.alignment = alignof(T);
180   request.usage = Allocation::Usage::Create;
181 
182   auto alloc = allocate(request);
183   new (alloc.ptr) T(std::forward<ARGS>(args)...);
184   return reinterpret_cast<T*>(alloc.ptr);
185 }
186 
187 template <typename T>
destroy(T * object)188 void Allocator::destroy(T* object) {
189   object->~T();
190 
191   Allocation alloc;
192   alloc.ptr = object;
193   alloc.request.size = sizeof(T);
194   alloc.request.alignment = alignof(T);
195   alloc.request.usage = Allocation::Usage::Create;
196   free(alloc);
197 }
198 
199 template <typename T, typename... ARGS>
make_unique(ARGS &&...args)200 Allocator::unique_ptr<T> Allocator::make_unique(ARGS&&... args) {
201   return make_unique_n<T>(1, std::forward<ARGS>(args)...);
202 }
203 
204 template <typename T, typename... ARGS>
make_unique_n(size_t n,ARGS &&...args)205 Allocator::unique_ptr<T> Allocator::make_unique_n(size_t n, ARGS&&... args) {
206   if (n == 0) {
207     return nullptr;
208   }
209 
210   Allocation::Request request;
211   request.size = sizeof(T) * n;
212   request.alignment = alignof(T);
213   request.usage = Allocation::Usage::Create;
214 
215   auto alloc = allocate(request);
216   new (alloc.ptr) T(std::forward<ARGS>(args)...);
217   return unique_ptr<T>(reinterpret_cast<T*>(alloc.ptr), Deleter{this, n});
218 }
219 
220 template <typename T, typename... ARGS>
make_shared(ARGS &&...args)221 std::shared_ptr<T> Allocator::make_shared(ARGS&&... args) {
222   Allocation::Request request;
223   request.size = sizeof(T);
224   request.alignment = alignof(T);
225   request.usage = Allocation::Usage::Create;
226 
227   auto alloc = allocate(request);
228   new (alloc.ptr) T(std::forward<ARGS>(args)...);
229   return std::shared_ptr<T>(reinterpret_cast<T*>(alloc.ptr), Deleter{this, 1});
230 }
231 
232 ///////////////////////////////////////////////////////////////////////////////
233 // TrackedAllocator
234 ///////////////////////////////////////////////////////////////////////////////
235 
236 // TrackedAllocator wraps an Allocator to track the allocations made.
237 class TrackedAllocator : public Allocator {
238  public:
239   struct UsageStats {
240     // Total number of allocations.
241     size_t count = 0;
242     // total allocation size in bytes (as requested, may be higher due to
243     // alignment or guards).
244     size_t bytes = 0;
245   };
246 
247   struct Stats {
248     // numAllocations() returns the total number of allocations across all
249     // usages for the allocator.
250     inline size_t numAllocations() const;
251 
252     // bytesAllocated() returns the total number of bytes allocated across all
253     // usages for the allocator.
254     inline size_t bytesAllocated() const;
255 
256     // Statistics per usage.
257     std::array<UsageStats, size_t(Allocation::Usage::Count)> byUsage;
258   };
259 
260   // Constructor that wraps an existing allocator.
261   inline TrackedAllocator(Allocator* allocator);
262 
263   // stats() returns the current allocator statistics.
264   inline Stats stats();
265 
266   // Allocator compliance
267   inline Allocation allocate(const Allocation::Request&) override;
268   inline void free(const Allocation&) override;
269 
270  private:
271   Allocator* const allocator;
272   std::mutex mutex;
273   Stats stats_;
274 };
275 
numAllocations()276 size_t TrackedAllocator::Stats::numAllocations() const {
277   size_t out = 0;
278   for (auto& stats : byUsage) {
279     out += stats.count;
280   }
281   return out;
282 }
283 
bytesAllocated()284 size_t TrackedAllocator::Stats::bytesAllocated() const {
285   size_t out = 0;
286   for (auto& stats : byUsage) {
287     out += stats.bytes;
288   }
289   return out;
290 }
291 
TrackedAllocator(Allocator * allocator)292 TrackedAllocator::TrackedAllocator(Allocator* allocator)
293     : allocator(allocator) {}
294 
stats()295 TrackedAllocator::Stats TrackedAllocator::stats() {
296   std::unique_lock<std::mutex> lock(mutex);
297   return stats_;
298 }
299 
allocate(const Allocation::Request & request)300 Allocation TrackedAllocator::allocate(const Allocation::Request& request) {
301   {
302     std::unique_lock<std::mutex> lock(mutex);
303     auto& usageStats = stats_.byUsage[int(request.usage)];
304     ++usageStats.count;
305     usageStats.bytes += request.size;
306   }
307   return allocator->allocate(request);
308 }
309 
free(const Allocation & allocation)310 void TrackedAllocator::free(const Allocation& allocation) {
311   {
312     std::unique_lock<std::mutex> lock(mutex);
313     auto& usageStats = stats_.byUsage[int(allocation.request.usage)];
314     MARL_ASSERT(usageStats.count > 0,
315                 "TrackedAllocator detected abnormal free()");
316     MARL_ASSERT(usageStats.bytes >= allocation.request.size,
317                 "TrackedAllocator detected abnormal free()");
318     --usageStats.count;
319     usageStats.bytes -= allocation.request.size;
320   }
321   return allocator->free(allocation);
322 }
323 
324 ///////////////////////////////////////////////////////////////////////////////
325 // StlAllocator
326 ///////////////////////////////////////////////////////////////////////////////
327 
328 // StlAllocator exposes an STL-compatible allocator wrapping a marl::Allocator.
329 template <typename T>
330 struct StlAllocator {
331   using value_type = T;
332   using pointer = T*;
333   using const_pointer = const T*;
334   using reference = T&;
335   using const_reference = const T&;
336   using size_type = size_t;
337   using difference_type = size_t;
338 
339   // An equivalent STL allocator for a different type.
340   template <class U>
341   struct rebind {
342     typedef StlAllocator<U> other;
343   };
344 
345   // Constructs an StlAllocator that will allocate using allocator.
346   // allocator must remain valid until this StlAllocator has been destroyed.
347   inline StlAllocator(Allocator* allocator);
348 
349   template <typename U>
350   inline StlAllocator(const StlAllocator<U>& other);
351 
352   // Returns the actual address of x even in presence of overloaded operator&.
353   inline pointer address(reference x) const;
354   inline const_pointer address(const_reference x) const;
355 
356   // Allocates the memory for n objects of type T.
357   // Does not actually construct the objects.
358   inline T* allocate(std::size_t n);
359 
360   // Deallocates the memory for n objects of type T.
361   inline void deallocate(T* p, std::size_t n);
362 
363   // Returns the maximum theoretically possible number of T stored in this
364   // allocator.
365   inline size_type max_size() const;
366 
367   // Copy constructs an object of type T at the address p.
368   inline void construct(pointer p, const_reference val);
369 
370   // Constructs an object of type U at the address P forwarning all other
371   // arguments to the constructor.
372   template <typename U, typename... Args>
373   inline void construct(U* p, Args&&... args);
374 
375   // Deconstructs the object at p. It does not free the memory.
376   inline void destroy(pointer p);
377 
378   // Deconstructs the object at p. It does not free the memory.
379   template <typename U>
380   inline void destroy(U* p);
381 
382  private:
383   inline Allocation::Request request(size_t n) const;
384 
385   template <typename U>
386   friend struct StlAllocator;
387   Allocator* allocator;
388 };
389 
390 template <typename T>
StlAllocator(Allocator * allocator)391 StlAllocator<T>::StlAllocator(Allocator* allocator) : allocator(allocator) {}
392 
393 template <typename T>
394 template <typename U>
StlAllocator(const StlAllocator<U> & other)395 StlAllocator<T>::StlAllocator(const StlAllocator<U>& other) {
396   allocator = other.allocator;
397 }
398 
399 template <typename T>
address(reference x)400 typename StlAllocator<T>::pointer StlAllocator<T>::address(reference x) const {
401   return &x;
402 }
403 template <typename T>
address(const_reference x)404 typename StlAllocator<T>::const_pointer StlAllocator<T>::address(
405     const_reference x) const {
406   return &x;
407 }
408 
409 template <typename T>
allocate(std::size_t n)410 T* StlAllocator<T>::allocate(std::size_t n) {
411   auto alloc = allocator->allocate(request(n));
412   return reinterpret_cast<T*>(alloc.ptr);
413 }
414 
415 template <typename T>
deallocate(T * p,std::size_t n)416 void StlAllocator<T>::deallocate(T* p, std::size_t n) {
417   Allocation alloc;
418   alloc.ptr = p;
419   alloc.request = request(n);
420   allocator->free(alloc);
421 }
422 
423 template <typename T>
max_size()424 typename StlAllocator<T>::size_type StlAllocator<T>::max_size() const {
425   return std::numeric_limits<size_type>::max() / sizeof(value_type);
426 }
427 
428 template <typename T>
construct(pointer p,const_reference val)429 void StlAllocator<T>::construct(pointer p, const_reference val) {
430   new (p) T(val);
431 }
432 
433 template <typename T>
434 template <typename U, typename... Args>
construct(U * p,Args &&...args)435 void StlAllocator<T>::construct(U* p, Args&&... args) {
436   ::new ((void*)p) U(std::forward<Args>(args)...);
437 }
438 
439 template <typename T>
destroy(pointer p)440 void StlAllocator<T>::destroy(pointer p) {
441   ((T*)p)->~T();
442 }
443 
444 template <typename T>
445 template <typename U>
destroy(U * p)446 void StlAllocator<T>::destroy(U* p) {
447   p->~U();
448 }
449 
450 template <typename T>
request(size_t n)451 Allocation::Request StlAllocator<T>::request(size_t n) const {
452   Allocation::Request req = {};
453   req.size = sizeof(T) * n;
454   req.alignment = alignof(T);
455   req.usage = Allocation::Usage::Stl;
456   return req;
457 }
458 
459 }  // namespace marl
460 
461 #endif  // marl_memory_h
462