1 | #pragma once |
2 | |
3 | #include <string.h> |
4 | #include <memory> |
5 | #include <vector> |
6 | #include <boost/noncopyable.hpp> |
7 | #include <common/likely.h> |
8 | #if __has_include(<sanitizer/asan_interface.h>) |
9 | # include <sanitizer/asan_interface.h> |
10 | #endif |
11 | #include <Core/Defines.h> |
12 | #include <Common/memcpySmall.h> |
13 | #include <Common/ProfileEvents.h> |
14 | #include <Common/Allocator.h> |
15 | |
16 | |
17 | namespace ProfileEvents |
18 | { |
19 | extern const Event ArenaAllocChunks; |
20 | extern const Event ArenaAllocBytes; |
21 | } |
22 | |
23 | namespace DB |
24 | { |
25 | |
26 | |
27 | /** Memory pool to append something. For example, short strings. |
28 | * Usage scenario: |
29 | * - put lot of strings inside pool, keep their addresses; |
30 | * - addresses remain valid during lifetime of pool; |
31 | * - at destruction of pool, all memory is freed; |
32 | * - memory is allocated and freed by large chunks; |
33 | * - freeing parts of data is not possible (but look at ArenaWithFreeLists if you need); |
34 | */ |
35 | class Arena : private boost::noncopyable |
36 | { |
37 | private: |
38 | /// Padding allows to use 'memcpySmallAllowReadWriteOverflow15' instead of 'memcpy'. |
39 | static constexpr size_t pad_right = 15; |
40 | |
41 | /// Contiguous chunk of memory and pointer to free space inside it. Member of single-linked list. |
42 | struct alignas(16) Chunk : private Allocator<false> /// empty base optimization |
43 | { |
44 | char * begin; |
45 | char * pos; |
46 | char * end; /// does not include padding. |
47 | |
48 | Chunk * prev; |
49 | |
50 | Chunk(size_t size_, Chunk * prev_) |
51 | { |
52 | ProfileEvents::increment(ProfileEvents::ArenaAllocChunks); |
53 | ProfileEvents::increment(ProfileEvents::ArenaAllocBytes, size_); |
54 | |
55 | begin = reinterpret_cast<char *>(Allocator<false>::alloc(size_)); |
56 | pos = begin; |
57 | end = begin + size_ - pad_right; |
58 | prev = prev_; |
59 | |
60 | ASAN_POISON_MEMORY_REGION(begin, size_); |
61 | } |
62 | |
63 | ~Chunk() |
64 | { |
65 | /// We must unpoison the memory before returning to the allocator, |
66 | /// because the allocator might not have asan integration, and the |
67 | /// memory would stay poisoned forever. If the allocator supports |
68 | /// asan, it will correctly poison the memory by itself. |
69 | ASAN_UNPOISON_MEMORY_REGION(begin, size()); |
70 | |
71 | Allocator<false>::free(begin, size()); |
72 | |
73 | if (prev) |
74 | delete prev; |
75 | } |
76 | |
77 | size_t size() const { return end + pad_right - begin; } |
78 | size_t remaining() const { return end - pos; } |
79 | }; |
80 | |
81 | size_t growth_factor; |
82 | size_t linear_growth_threshold; |
83 | |
84 | /// Last contiguous chunk of memory. |
85 | Chunk * head; |
86 | size_t size_in_bytes; |
87 | |
88 | static size_t roundUpToPageSize(size_t s) |
89 | { |
90 | return (s + 4096 - 1) / 4096 * 4096; |
91 | } |
92 | |
93 | /// If chunks size is less than 'linear_growth_threshold', then use exponential growth, otherwise - linear growth |
94 | /// (to not allocate too much excessive memory). |
95 | size_t nextSize(size_t min_next_size) const |
96 | { |
97 | size_t size_after_grow = 0; |
98 | |
99 | if (head->size() < linear_growth_threshold) |
100 | { |
101 | size_after_grow = std::max(min_next_size, head->size() * growth_factor); |
102 | } |
103 | else |
104 | { |
105 | // allocContinue() combined with linear growth results in quadratic |
106 | // behavior: we append the data by small amounts, and when it |
107 | // doesn't fit, we create a new chunk and copy all the previous data |
108 | // into it. The number of times we do this is directly proportional |
109 | // to the total size of data that is going to be serialized. To make |
110 | // the copying happen less often, round the next size up to the |
111 | // linear_growth_threshold. |
112 | size_after_grow = ((min_next_size + linear_growth_threshold - 1) |
113 | / linear_growth_threshold) * linear_growth_threshold; |
114 | } |
115 | |
116 | assert(size_after_grow >= min_next_size); |
117 | return roundUpToPageSize(size_after_grow); |
118 | } |
119 | |
120 | /// Add next contiguous chunk of memory with size not less than specified. |
121 | void NO_INLINE addChunk(size_t min_size) |
122 | { |
123 | head = new Chunk(nextSize(min_size + pad_right), head); |
124 | size_in_bytes += head->size(); |
125 | } |
126 | |
127 | friend class ArenaAllocator; |
128 | template <size_t> friend class AlignedArenaAllocator; |
129 | |
130 | public: |
131 | Arena(size_t initial_size_ = 4096, size_t growth_factor_ = 2, size_t linear_growth_threshold_ = 128 * 1024 * 1024) |
132 | : growth_factor(growth_factor_), linear_growth_threshold(linear_growth_threshold_), |
133 | head(new Chunk(initial_size_, nullptr)), size_in_bytes(head->size()) |
134 | { |
135 | } |
136 | |
137 | ~Arena() |
138 | { |
139 | delete head; |
140 | } |
141 | |
142 | /// Get piece of memory, without alignment. |
143 | char * alloc(size_t size) |
144 | { |
145 | if (unlikely(head->pos + size > head->end)) |
146 | addChunk(size); |
147 | |
148 | char * res = head->pos; |
149 | head->pos += size; |
150 | ASAN_UNPOISON_MEMORY_REGION(res, size + pad_right); |
151 | return res; |
152 | } |
153 | |
154 | /// Get peice of memory with alignment |
155 | char * alignedAlloc(size_t size, size_t alignment) |
156 | { |
157 | do |
158 | { |
159 | void * head_pos = head->pos; |
160 | size_t space = head->end - head->pos; |
161 | |
162 | auto res = static_cast<char *>(std::align(alignment, size, head_pos, space)); |
163 | if (res) |
164 | { |
165 | head->pos = static_cast<char *>(head_pos); |
166 | head->pos += size; |
167 | ASAN_UNPOISON_MEMORY_REGION(res, size + pad_right); |
168 | return res; |
169 | } |
170 | |
171 | addChunk(size + alignment); |
172 | } while (true); |
173 | } |
174 | |
175 | template <typename T> |
176 | T * alloc() |
177 | { |
178 | return reinterpret_cast<T *>(alignedAlloc(sizeof(T), alignof(T))); |
179 | } |
180 | |
181 | /** Rollback just performed allocation. |
182 | * Must pass size not more that was just allocated. |
183 | * Return the resulting head pointer, so that the caller can assert that |
184 | * the allocation it intended to roll back was indeed the last one. |
185 | */ |
186 | void * rollback(size_t size) |
187 | { |
188 | head->pos -= size; |
189 | ASAN_POISON_MEMORY_REGION(head->pos, size + pad_right); |
190 | return head->pos; |
191 | } |
192 | |
193 | /** Begin or expand a contiguous range of memory. |
194 | * 'range_start' is the start of range. If nullptr, a new range is |
195 | * allocated. |
196 | * If there is no space in the current chunk to expand the range, |
197 | * the entire range is copied to a new, bigger memory chunk, and the value |
198 | * of 'range_start' is updated. |
199 | * If the optional 'start_alignment' is specified, the start of range is |
200 | * kept aligned to this value. |
201 | * |
202 | * NOTE This method is usable only for the last allocation made on this |
203 | * Arena. For earlier allocations, see 'realloc' method. |
204 | */ |
205 | char * allocContinue(size_t additional_bytes, char const *& range_start, |
206 | size_t start_alignment = 0) |
207 | { |
208 | if (!range_start) |
209 | { |
210 | // Start a new memory range. |
211 | char * result = start_alignment |
212 | ? alignedAlloc(additional_bytes, start_alignment) |
213 | : alloc(additional_bytes); |
214 | |
215 | range_start = result; |
216 | return result; |
217 | } |
218 | |
219 | // Extend an existing memory range with 'additional_bytes'. |
220 | |
221 | // This method only works for extending the last allocation. For lack of |
222 | // original size, check a weaker condition: that 'begin' is at least in |
223 | // the current Chunk. |
224 | assert(range_start >= head->begin && range_start < head->end); |
225 | |
226 | if (head->pos + additional_bytes <= head->end) |
227 | { |
228 | // The new size fits into the last chunk, so just alloc the |
229 | // additional size. We can alloc without alignment here, because it |
230 | // only applies to the start of the range, and we don't change it. |
231 | return alloc(additional_bytes); |
232 | } |
233 | |
234 | // New range doesn't fit into this chunk, will copy to a new one. |
235 | // |
236 | // Note: among other things, this method is used to provide a hack-ish |
237 | // implementation of realloc over Arenas in ArenaAllocators. It wastes a |
238 | // lot of memory -- quadratically so when we reach the linear allocation |
239 | // threshold. This deficiency is intentionally left as is, and should be |
240 | // solved not by complicating this method, but by rethinking the |
241 | // approach to memory management for aggregate function states, so that |
242 | // we can provide a proper realloc(). |
243 | const size_t existing_bytes = head->pos - range_start; |
244 | const size_t new_bytes = existing_bytes + additional_bytes; |
245 | const char * old_range = range_start; |
246 | |
247 | char * new_range = start_alignment |
248 | ? alignedAlloc(new_bytes, start_alignment) |
249 | : alloc(new_bytes); |
250 | |
251 | memcpy(new_range, old_range, existing_bytes); |
252 | |
253 | range_start = new_range; |
254 | return new_range + existing_bytes; |
255 | } |
256 | |
257 | /// NOTE Old memory region is wasted. |
258 | char * realloc(const char * old_data, size_t old_size, size_t new_size) |
259 | { |
260 | char * res = alloc(new_size); |
261 | if (old_data) |
262 | { |
263 | memcpy(res, old_data, old_size); |
264 | ASAN_POISON_MEMORY_REGION(old_data, old_size); |
265 | } |
266 | return res; |
267 | } |
268 | |
269 | char * alignedRealloc(const char * old_data, size_t old_size, size_t new_size, size_t alignment) |
270 | { |
271 | char * res = alignedAlloc(new_size, alignment); |
272 | if (old_data) |
273 | { |
274 | memcpy(res, old_data, old_size); |
275 | ASAN_POISON_MEMORY_REGION(old_data, old_size); |
276 | } |
277 | return res; |
278 | } |
279 | |
280 | /// Insert string without alignment. |
281 | const char * insert(const char * data, size_t size) |
282 | { |
283 | char * res = alloc(size); |
284 | memcpy(res, data, size); |
285 | return res; |
286 | } |
287 | |
288 | const char * alignedInsert(const char * data, size_t size, size_t alignment) |
289 | { |
290 | char * res = alignedAlloc(size, alignment); |
291 | memcpy(res, data, size); |
292 | return res; |
293 | } |
294 | |
295 | /// Size of chunks in bytes. |
296 | size_t size() const |
297 | { |
298 | return size_in_bytes; |
299 | } |
300 | |
301 | size_t remainingSpaceInCurrentChunk() const |
302 | { |
303 | return head->remaining(); |
304 | } |
305 | }; |
306 | |
307 | using ArenaPtr = std::shared_ptr<Arena>; |
308 | using Arenas = std::vector<ArenaPtr>; |
309 | |
310 | |
311 | } |
312 | |