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
20 #include <stdint.h>
21
22 #include <array>
23 #include <cstdlib>
24 #include <memory>
25 #include <mutex>
26 #include <utility> // std::forward
27
28 namespace marl {
29
30 // pageSize() returns the size in bytes of a virtual memory page for the host
31 // system.
32 size_t pageSize();
33
34 // Allocation holds the result of a memory allocation from an Allocator.
35 struct Allocation {
36 // Intended usage of the allocation. Used for allocation trackers.
37 enum class Usage {
38 Undefined = 0,
39 Stack, // Fiber stack
40 Create, // Allocator::create(), make_unique(), make_shared()
41 Vector, // marl::vector<T>
42 Count, // Not intended to be used as a usage type - used for upper bound.
43 };
44
45 // Request holds all the information required to make an allocation.
46 struct Request {
47 size_t size = 0; // The size of the allocation in bytes.
48 size_t alignment = 0; // The minimum alignment of the allocation.
49 bool useGuards = false; // Whether the allocation is guarded.
50 Usage usage = Usage::Undefined; // Intended usage of the allocation.
51 };
52
53 void* ptr = nullptr; // The pointer to the allocated memory.
54 Request request; // Request used for the allocation.
55 };
56
57 // Allocator is an interface to a memory allocator.
58 // Marl provides a default implementation with Allocator::Default.
59 class Allocator {
60 public:
61 // The default allocator. Initialized with an implementation that allocates
62 // from the OS. Can be assigned a custom implementation.
63 static Allocator* Default;
64
65 // Deleter is a smart-pointer compatible deleter that can be used to delete
66 // objects created by Allocator::create(). Deleter is used by the smart
67 // pointers returned by make_shared() and make_unique().
68 struct Deleter {
69 inline Deleter();
70 inline Deleter(Allocator* allocator);
71
72 template <typename T>
73 inline void operator()(T* object);
74
75 Allocator* allocator = nullptr;
76 };
77
78 // unique_ptr<T> is an alias to std::unique_ptr<T, Deleter>.
79 template <typename T>
80 using unique_ptr = std::unique_ptr<T, Deleter>;
81
82 virtual ~Allocator() = default;
83
84 // allocate() allocates memory from the allocator.
85 // The returned Allocation::request field must be equal to the Request
86 // parameter.
87 virtual Allocation allocate(const Allocation::Request&) = 0;
88
89 // free() frees the memory returned by allocate().
90 // The Allocation must have all fields equal to those returned by allocate().
91 virtual void free(const Allocation&) = 0;
92
93 // create() allocates and constructs an object of type T, respecting the
94 // alignment of the type.
95 // The pointer returned by create() must be deleted with destroy().
96 template <typename T, typename... ARGS>
97 inline T* create(ARGS&&... args);
98
99 // destroy() destructs and frees the object allocated with create().
100 template <typename T>
101 inline void destroy(T* object);
102
103 // make_unique() returns a new object allocated from the allocator wrapped
104 // in a unique_ptr that respects the alignemnt of the type.
105 template <typename T, typename... ARGS>
106 inline unique_ptr<T> make_unique(ARGS&&... args);
107
108 // make_shared() returns a new object allocated from the allocator
109 // wrapped in a std::shared_ptr that respects the alignemnt of the type.
110 template <typename T, typename... ARGS>
111 inline std::shared_ptr<T> make_shared(ARGS&&... args);
112
113 protected:
114 Allocator() = default;
115 };
116
Deleter()117 Allocator::Deleter::Deleter() : allocator(nullptr) {}
Deleter(Allocator * allocator)118 Allocator::Deleter::Deleter(Allocator* allocator) : allocator(allocator) {}
119
120 template <typename T>
operator()121 void Allocator::Deleter::operator()(T* object) {
122 object->~T();
123
124 Allocation allocation;
125 allocation.ptr = object;
126 allocation.request.size = sizeof(T);
127 allocation.request.alignment = alignof(T);
128 allocation.request.usage = Allocation::Usage::Create;
129 allocator->free(allocation);
130 }
131
132 template <typename T, typename... ARGS>
create(ARGS &&...args)133 T* Allocator::create(ARGS&&... args) {
134 Allocation::Request request;
135 request.size = sizeof(T);
136 request.alignment = alignof(T);
137 request.usage = Allocation::Usage::Create;
138
139 auto alloc = allocate(request);
140 new (alloc.ptr) T(std::forward<ARGS>(args)...);
141 return reinterpret_cast<T*>(alloc.ptr);
142 }
143
144 template <typename T>
destroy(T * object)145 void Allocator::destroy(T* object) {
146 object->~T();
147
148 Allocation alloc;
149 alloc.ptr = object;
150 alloc.request.size = sizeof(T);
151 alloc.request.alignment = alignof(T);
152 alloc.request.usage = Allocation::Usage::Create;
153 free(alloc);
154 }
155
156 template <typename T, typename... ARGS>
make_unique(ARGS &&...args)157 Allocator::unique_ptr<T> Allocator::make_unique(ARGS&&... args) {
158 Allocation::Request request;
159 request.size = sizeof(T);
160 request.alignment = alignof(T);
161 request.usage = Allocation::Usage::Create;
162
163 auto alloc = allocate(request);
164 new (alloc.ptr) T(std::forward<ARGS>(args)...);
165 return unique_ptr<T>(reinterpret_cast<T*>(alloc.ptr), Deleter{this});
166 }
167
168 template <typename T, typename... ARGS>
make_shared(ARGS &&...args)169 std::shared_ptr<T> Allocator::make_shared(ARGS&&... args) {
170 Allocation::Request request;
171 request.size = sizeof(T);
172 request.alignment = alignof(T);
173 request.usage = Allocation::Usage::Create;
174
175 auto alloc = allocate(request);
176 new (alloc.ptr) T(std::forward<ARGS>(args)...);
177 return std::shared_ptr<T>(reinterpret_cast<T*>(alloc.ptr), Deleter{this});
178 }
179
180 // aligned_storage() is a replacement for std::aligned_storage that isn't busted
181 // on older versions of MSVC.
182 template <size_t SIZE, size_t ALIGNMENT>
183 struct aligned_storage {
184 struct alignas(ALIGNMENT) type {
185 unsigned char data[SIZE];
186 };
187 };
188
189 // TrackedAllocator wraps an Allocator to track the allocations made.
190 class TrackedAllocator : public Allocator {
191 public:
192 struct UsageStats {
193 // Total number of allocations.
194 size_t count = 0;
195 // total allocation size in bytes (as requested, may be higher due to
196 // alignment or guards).
197 size_t bytes = 0;
198 };
199
200 struct Stats {
201 // numAllocations() returns the total number of allocations across all
202 // usages for the allocator.
203 inline size_t numAllocations() const;
204
205 // bytesAllocated() returns the total number of bytes allocated across all
206 // usages for the allocator.
207 inline size_t bytesAllocated() const;
208
209 // Statistics per usage.
210 std::array<UsageStats, size_t(Allocation::Usage::Count)> byUsage;
211 };
212
213 // Constructor that wraps an existing allocator.
214 inline TrackedAllocator(Allocator* allocator);
215
216 // stats() returns the current allocator statistics.
217 inline Stats stats();
218
219 // Allocator compliance
220 inline Allocation allocate(const Allocation::Request&) override;
221 inline void free(const Allocation&) override;
222
223 private:
224 Allocator* const allocator;
225 std::mutex mutex;
226 Stats stats_;
227 };
228
numAllocations()229 size_t TrackedAllocator::Stats::numAllocations() const {
230 size_t out = 0;
231 for (auto& stats : byUsage) {
232 out += stats.count;
233 }
234 return out;
235 }
236
bytesAllocated()237 size_t TrackedAllocator::Stats::bytesAllocated() const {
238 size_t out = 0;
239 for (auto& stats : byUsage) {
240 out += stats.bytes;
241 }
242 return out;
243 }
244
TrackedAllocator(Allocator * allocator)245 TrackedAllocator::TrackedAllocator(Allocator* allocator)
246 : allocator(allocator) {}
247
stats()248 TrackedAllocator::Stats TrackedAllocator::stats() {
249 std::unique_lock<std::mutex> lock(mutex);
250 return stats_;
251 }
252
allocate(const Allocation::Request & request)253 Allocation TrackedAllocator::allocate(const Allocation::Request& request) {
254 {
255 std::unique_lock<std::mutex> lock(mutex);
256 auto& usageStats = stats_.byUsage[int(request.usage)];
257 ++usageStats.count;
258 usageStats.bytes += request.size;
259 }
260 return allocator->allocate(request);
261 }
262
free(const Allocation & allocation)263 void TrackedAllocator::free(const Allocation& allocation) {
264 {
265 std::unique_lock<std::mutex> lock(mutex);
266 auto& usageStats = stats_.byUsage[int(allocation.request.usage)];
267 MARL_ASSERT(usageStats.count > 0,
268 "TrackedAllocator detected abnormal free()");
269 MARL_ASSERT(usageStats.bytes >= allocation.request.size,
270 "TrackedAllocator detected abnormal free()");
271 --usageStats.count;
272 usageStats.bytes -= allocation.request.size;
273 }
274 return allocator->free(allocation);
275 }
276
277 } // namespace marl
278
279 #endif // marl_memory_h
280