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
17namespace ProfileEvents
18{
19 extern const Event ArenaAllocChunks;
20 extern const Event ArenaAllocBytes;
21}
22
23namespace 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 */
35class Arena : private boost::noncopyable
36{
37private:
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
130public:
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
307using ArenaPtr = std::shared_ptr<Arena>;
308using Arenas = std::vector<ArenaPtr>;
309
310
311}
312