1/* @title: Memory Arena Provider API */
2#pragma once
3#include <math/bit.h>
4#include <mem/arena.h>
5#include <structures/bitmap.h>
6#include <structures/hlist.h>
7#include <structures/mpmc_lifo_list.h>
8#include <structures/rbt.h>
9#include <sync/mutex.h>
10#include <sync/once_token.h>
11#include <sync/rcu.h>
12#include <sync/rwlock.h>
13#include <types/types.h>
14
15/* TODO: Arena unregistration */
16
17struct log_site;
18struct arena_budget;
19struct arena_dumpster;
20
21/* For arena_fn_desc, which is the default set of
22 * functions. This primarily serves to do a little
23 * bit of debugging sanitization/warning, such as
24 * identifying overloaded functions and properly
25 * checking if the provided *p is stack allocated. Also
26 * works with arena_caps to keep these bits in the bottom dword */
27#define ARENA_MAX_FN 32
28
29#define ARENA_MAX_SEG 16 /* 4 bits are used to store segment IDs */
30#define ARENA_MAX_EXT_FN 16
31
32#define ARENA_SEG_CUSTOM(s) ((s) + ARENA_SEG_MAX)
33#define ARENA_SEG_ALIGN 8u
34#define ARENA_SEG_MAX_SIZE 4096
35
36#define ARENA_BIN_NOT_HEAD_MAGIC 0xB0D1ED100B0D1ED1 /* "BODIED! BODIED!" */
37
38/* These arena segments are the IDs for segments[] in struct arena, so that
39 * various subsystems can extend the structure and get extra storage space.
40 *
41 * Note that these names are NOT the *only* uses of their respective
42 * IDs. ID 0 can very well be whatever the arena provider wants,
43 * however, this enum is more of providing "norms" for using the segments
44 */
45enum arena_seg_id : uint16_t {
46 ARENA_SEG_PROPS = 0,
47 ARENA_SEG_SYNC = 1,
48 ARENA_SEG_LOG = 2,
49 ARENA_SEG_TRACKING = 3,
50 ARENA_SEG_BUDGET = 4,
51 ARENA_SEG_DEBUG = 5,
52 ARENA_SEG_MAX, /* Maximum for the default segment types */
53};
54static_assert(ARENA_SEG_MAX <= ARENA_MAX_SEG);
55
56/* PARAMS indicates that *p is actually
57 * meant to be a stack allocated set of parameters */
58enum arena_fn_type { ARENA_FN_PTR, ARENA_FN_PARAMS };
59
60enum arena_hook_type {
61 ARENA_HOOK_RECYCLE,
62 ARENA_HOOK_GC_ENTER,
63};
64
65enum arena_desc_flags {
66 /* All arenas of this descriptor support garbage collection.
67 * Destruction will enqueue it to GC */
68 ARENA_DESC_GC = 1,
69};
70
71/* arena_caps: 64 bit bitflags
72 *
73 * ┌──────────────────────────────────────────┐
74 * Bits │ 63..60 59..56 55..52 51..48 47..32 31..0 │
75 * Use │ AAAA AAAA AAAA AAAA EEEE FFFF │
76 * └──────────────────────────────────────────┘
77 *
78 * F - Default arena_ops
79 * E - Extended arena_ops
80 * A - Unused (available)
81 * * - Unused (unavailable)
82 *
83 */
84enum arena_caps { ARENA_CAP };
85
86/* Notes on prefixes, since this op table gets busy
87 *
88 * qry_* is the per-object/ptr "getter", set_* for the "setters", and
89 * get_* is more of an "extract this from the structure", and
90 * getref/putref are the reference counting ops
91 *
92 * As for arena_params, we only use this whenever no *p exists. In other cases,
93 * *p is intended to be used as the arena allocator as an overloadable
94 * field that can potentially point to a struct arena_params in some
95 * implementations.
96 */
97struct arena_ops {
98 struct arena *(*create)(struct arena_params *params);
99 void *(*alloc)(struct arena *a, size_t size, struct alloc_params params);
100
101 void *(*realloc)(struct arena *a, void *ptr, size_t size,
102 struct alloc_params params);
103
104 /* Just in case the arena wants to have extra functionality
105 * on allocation. Not an ext_fn since it's got so many arguments,
106 * and exists at all because alloc does not take in a *p. */
107 void *(*alloc_special)(struct arena *a, struct arena_params *params,
108 struct alloc_params a_params);
109
110 /* If an arena cannot implement a free operation without other
111 * information, it should be using an extended function */
112 enum err (*free)(struct arena *a, void *p, struct alloc_params params);
113
114 /* Aligned variants */
115 void *(*alloc_aligned)(struct arena *a, size_t size, size_t align,
116 struct alloc_params params);
117
118 enum err (*free_aligned)(struct arena *a, void *p, size_t align,
119 struct alloc_params params);
120
121 /* Responsible for clearing out all memory allocations
122 * inside of the arena. It does not necessarily reset its
123 * internal structures */
124 enum err (*clear)(struct arena *a, struct arena_params *params);
125
126 /*
127 * Does NOT ever free *a itself. This is the raw internal
128 * free function, and should consider assertion that the
129 * refcount == 0 if reference counting is being used
130 */
131 void (*destroy)(struct arena *a, struct arena_params *params);
132
133 /* Handles actually resetting internal tracking data structures to
134 * a point where it is effectively identical to creation time.
135 *
136 * Does NOT clear(), and can panic
137 * if called without a cleared arena */
138 void (*reset)(struct arena *a, struct arena_params *params);
139
140 enum err (*set_tag)(struct arena *a, void *p, arena_tag_t tag);
141
142 enum err (*set_budget)(struct arena *a, struct arena_budget b);
143
144 size_t (*qry_size)(struct arena *a, void *p);
145
146 arena_tag_t (*qry_tag)(struct arena *a, void *p);
147
148 stack_handle_t (*qry_trace)(struct arena *a, void *p);
149
150 bool (*qry_owns)(struct arena *a, void *p);
151
152 /* *params is for the outside world */
153 struct log_site *(*get_log_site)(struct arena *a,
154 struct arena_params *params);
155
156 struct arena_budget (*get_budget)(struct arena *a,
157 struct arena_params *params);
158
159 struct arena_dumpster *(*get_dumpster)(struct arena *a,
160 struct arena_params *params);
161
162 arena_identity_t (*get_identity)(struct arena *a,
163 struct arena_params *params);
164
165 bool (*identity_eq)(struct arena *a, struct arena *b);
166
167 /* We do not have a on_gc_destruction. GC destruction
168 * calls destroy(), and that is our hook */
169 enum err (*on_hook)(struct arena *a, enum arena_hook_type hook);
170
171 bool (*is_empty)(struct arena *a);
172
173 /* These merely handle the internal reference counting, notably
174 * NOT destroying or recycling arenas */
175 bool (*getref)(struct arena *a);
176
177 void (*putref)(struct arena *a);
178
179 arena_ext_fn_t ext_fns[ARENA_MAX_EXT_FN];
180};
181static_assert((sizeof(struct arena_ops) -
182 sizeof(arena_ext_fn_t) * ARENA_MAX_EXT_FN) <
183 sizeof(arena_ext_fn_t) * ARENA_MAX_FN);
184
185#define ARENA_OP_OFFSET(op) (offsetof(struct arena_ops, op) / sizeof(any_fn_t))
186
187/* There's a 3-level hierarchy here:
188 *
189 * Bins store lockless singly linked lists of same-vsize arenas, and can be
190 * created on demand for rapid arena reuse
191 *
192 * Bins must point to dumpsters, which are caches of various sized bins,
193 * and control the "destroy" operation through a single detach operation
194 *
195 * Dumpsters can also be created on demand, and all dumpsters are tracked
196 * with a landfill. For a given type of arena, there can only be a single
197 * landfill globally, as its job is primarily profiling and tracking
198 * dumpsters so that on, say, an OOM, dumpsters and bins can be reached.
199 *
200 * ┌──────────────────────────────────────┐
201 * │ Arena Descriptor │
202 * └──────────────────────────────────────┘
203 * │
204 * ▼
205 * ┌──────────────────────────────────────┐
206 * │ Arena Landfill (only 1) │
207 * └──────────────────────────────────────┘
208 * │ │
209 * ▼ ▼
210 * ┌──────────────────┐┌──────────────────┐
211 * │ Arena Dumpster 1 ││ Arena Dumpster 2 │
212 * └──────────────────┘└──────────────────┘
213 * │
214 * ┌───┴───────────────────────┐
215 * ▼ ▼
216 * ┌──────────┐ ┌──────────┐
217 * │ Bucket 1 │ ● ● ● ● │ Bucket N │
218 * └──────────┘ └──────────┘
219 * │
220 * │
221 * ┌─────┘
222 * │ ┌───────────────┐ ┌───────────────┐
223 * └─▶│ Bin ID A Lead │──▶│ Bin ID B Lead │
224 * └───────────────┘ └───────────────┘
225 * │
226 * │
227 * ┌──────────┘
228 * │ ┌──────────────┐ ┌──────────────┐
229 * └─▶│ Bin ID A 1 │────▶│ Bin ID A 2 │
230 * └──────────────┘ └──────────────┘
231 * │
232 * │
233 * ┌──────────┘
234 * │ ┌─────────┐ ┌─────────┐ ┌─────────┐
235 * └─▶│ Arena 1 │─▶│ Arena 2 │─▶│ Arena 3 │
236 * └─────────┘ └─────────┘ └─────────┘
237 *
238 * The way that the dumpsters store bins is they select a single "leader", and
239 * subsequent followers. Bins can choose to not participate in the dumpster,
240 * and such bins would be independently reference counted.
241 */
242struct arena_bin {
243 union {
244 struct hlist_node hlist_node;
245 uintptr_t not_head_magic;
246 };
247
248 struct list_node bin_list;
249
250 atomic_bool dying;
251
252 uint32_t arena_size;
253
254 arena_identity_t identity;
255 refcount_t refcount;
256
257 struct mpmc_lifo_list arenas;
258 struct arena_dumpster *dumpster;
259
260 struct rcu_cb rcu;
261};
262
263/* Arena dumpster buckets only hold "leader" bins, which link
264 * identical identity bins under themselves. Upon arena bin deletion,
265 * the next-in-line bin is promoted to be the leader, so
266 * any dying bin that is read by an RCU reader under a bucket merely
267 * requires ->next traversal to find the successor leader */
268struct arena_bucket {
269 struct mutex mutex;
270 struct hlist_head bins;
271};
272
273struct arena_dumpster {
274 struct arena_landfill *landfill;
275
276 refcount_t refcount;
277 size_t bucket_count; /* Stays constant after init */
278 struct arena_bucket *buckets;
279};
280
281struct arena_landfill {
282 struct list_head dumpster_list;
283};
284
285/* NOTE: certain capabilities do not necessarily need to be implemented
286 * by the arena. For instance, if a zeroed allocation is requested but the
287 * arena has no zero allocation support, then we simply memset before
288 * giving the memory back */
289struct arena_strategy_capabilities {
290 struct alloc_capabilities alloc_caps;
291 enum arena_flags supported_flags;
292};
293
294struct arena_fn_desc {
295 any_fn_t fn;
296 enum arena_fn_type type;
297};
298
299struct arena_budget {
300 ALLOC_PRIORITY_BITMAP_DECLARE(allowed_prios);
301 sz_b_t used_bytes;
302 sz_b_t capacity_bytes;
303 void *data;
304};
305
306struct arena_seg_log {
307 struct log_site *site;
308 struct log_handle handle;
309};
310
311struct arena_seg_sync {
312 union {
313 struct spinlock spinlock;
314 struct mutex mutex;
315 struct rwlock rwlock;
316 };
317
318 refcount_t refcount;
319};
320
321struct arena_seg_props {
322 const char *name;
323 enum arena_flags flags;
324 size_t capacity;
325};
326
327/* Given back to callers */
328struct arena_seg {
329 uint16_t id;
330 uint8_t *storage;
331};
332
333struct arena_seg_desc {
334 uint16_t large : 1; /* Only one large segment per arena is allowed */
335 uint16_t id : BITS_NEEDED(ARENA_MAX_SEG - 1);
336
337 /* This is used for generic accessors that wrap around our provided
338 * arena_seg types, and when in DEBUG_ARENA, stored in the inmem_desc */
339 uint16_t type : BITS_NEEDED(ARENA_SEG_MAX - 1);
340 uint16_t size;
341};
342
343/* inmem_desc is for what actually resides in payload[], and it's
344 * a shrunken down, internal representation of segments that only
345 * holds extra data when debugging
346 *
347 * NOTE: Segment *indices* are completely undefined.
348 * What remains stable is the ID of the segment (0 <= id < 16).
349 * Segment lookups key based off the ID,
350 * and existence is only enforced
351 * in debug builds to keep the arena payload[] slim
352 */
353struct arena_seg_inmem_desc {
354 uint16_t id : 4;
355
356 /* Where is this segment, in the payload[]? */
357 uint16_t offset_bump : BITS_NEEDED(ARENA_SEG_MAX_SIZE - 1);
358
359#ifdef DEBUG_ARENA
360 uint16_t type : BITS_NEEDED(ARENA_SEG_MAX - 1);
361 uint16_t present : 1;
362 uint16_t size;
363 struct arena_seg *seg;
364#endif
365};
366
367static_assert(BITS_NEEDED(ARENA_MAX_SEG - 1) == 4);
368
369#ifndef DEBUG_ARENA
370ct_assert_struct_size_eq(arena_seg_inmem_desc, 2);
371#endif
372
373/* This is the parent structure of the operations and other
374 * per-arena implementation data */
375struct arena_desc {
376 struct list_node list_node;
377 const char *name;
378 const char *description;
379
380 enum arena_strategy strategy;
381 struct arena_strategy_capabilities caps;
382 struct arena_ops *ops;
383 struct arena_seg_desc seg_descs[ARENA_MAX_SEG];
384 struct arena_fn_desc fn_descs[ARENA_MAX_FN];
385 struct arena_dumpster *dumpster;
386};
387
388struct arena_desc_linker_record {
389 struct arena_desc *desc;
390};
391
392/* We keep arenas minimal: some arenas
393 * that might just be wrappers around things have no reason to
394 * take up any more than a qword. We do keep a *desc instead of, say,
395 * an enum arena_policy because descriptors can be ad-hoc created */
396struct arena {
397 union {
398 struct arena_desc *desc;
399
400 /* We reinterpret this qword as a list_node when the arena
401 * is being recycled and resides inside a bin. The reason
402 * this is safe is because the GC structures explicitly track
403 * their arena_desc objects, so we can easily get the
404 * *desc pointer back for this arena */
405 struct mpmc_lifo_list_node list_node;
406 };
407
408#ifdef DEBUG_ARENA
409 struct alloc_capabilities alloc_caps;
410 enum arena_flags flags;
411 uint16_t seg_count;
412 size_t total_size;
413 ONCE_TOKEN_DEFINE(struct thread *,
414 entered); /* Debug single threaded arenas */
415#endif
416
417 /*
418 * Every segment gets an inmem_desc for itself. Since inmem_descs are 16
419 * bits, having the offset_bump be absolute significantly limits the amount
420 * of storage for segments various arenas use.
421 * e.g.
422 *
423 * with
424 *
425 * [ inmem_desc id 1 offset_bump 8 ]
426 * [ inmem_desc id 2 offset_bump 24 ]
427 * [ inmem_desc id 3 offset_bump 4 ]
428 * [ inmem_desc id 4 offset_bump 2 ]
429 *
430 * traversing to id 1 is payload + 8, id 2 is payload + 8 + 24
431 * id 3 is payload + 8 + 24 + 4, and id 4 is payload + 8 + 24 + 4 + 2
432 *
433 * Thus, because we are using this form of cursor bumping,
434 * the largest a single segment can be is limited by its own size,
435 * and never the collective sizes of all segments.
436 *
437 * We use this idea of "cursor bumping" instead of a raw size because of
438 * the leading inmem_descs. If we only tracked sizes, we'd need to use more
439 * data to track the number of segments (unneeded information) in order
440 * to compute the offset of the first segment.
441 *
442 * Tracking raw sizes also makes the last inmem_desc's
443 * size redundant (it trails, nothing needs to know its size).
444 *
445 * It is permitted to have one very large segment, however,
446 * that must be explicitly flagged as trailing.
447 *
448 * We can still track the amount of segments
449 * by going from payload[0] up until the first segment (the delta being given by
450 * the cursor bump of the first in memory descriptor), and counting the amount
451 * of non-zero inmem_descs (zero being invalid due to size zero).
452 *
453 * ┌────────────────────────────────────────────────────────────┐
454 * │ struct arena's payload[] │
455 * └────────────────────────────────────────────────────────────┘
456 *
457 * │ │
458 * │ │
459 * │ │
460 * │ │
461 * ▼ ▼
462 *
463 * ┌──────────────┐ ┌──────────────┐ ┌──────────────────┐
464 * │ inmem_desc 0 │ │ inmem_desc 1 │ ● ● ● │descriptor storage│
465 * └──────────────┘ └──────────────┘ └──────────────────┘
466 * ││
467 * │└────────────────┬─────────────────┐
468 * │ │ │
469 * ▼ ▼ ▼
470 * ┌──────────────┐ ┌──────────────┐ ┌─────────────────┐
471 * │ segment ID │ │ payload[] │ │ other tracking │
472 * │ 4 bits │ │ cursor bump │ │ metadata │
473 * └──────────────┘ └──────────────┘ └─────────────────┘
474 */
475 cc_align_as(max_align_t) uint8_t payload[];
476};
477
478/* A wrapper around e^(x - n) where n is scale, and x is used */
479err_checked
480arena_budget_prio_scale(sz_b_t used, int scale,
481 ALLOC_PRIORITY_BITMAP_DECLARE(prio_map_out));
482
483/* Arena strategies call into this with their fully formed descriptors */
484struct arena *arena_create_full(struct arena_seg_desc *seg_descs,
485 size_t seg_count);
486struct arena_seg arena_seg_lookup(struct arena *a, uint16_t seg_id);
487struct arena_seg arena_seg_for_idx(struct arena *a, uint16_t idx);
488err_checked arena_desc_register(struct arena_desc *d);
489struct arena_desc *arena_desc_lookup(enum arena_strategy strat);
490
491/* NOTE: with no segments, this is UB */
492size_t arena_seg_count(struct arena *a);
493
494void arena_dumpster_add_bin(struct arena_dumpster *dumpster,
495 struct arena_bin *bin);
496
497#define arena_for_each_seg(seg_, arena_) \
498 for (uint16_t __seg_idx = 0; \
499 __seg_idx < arena_seg_count(arena_) && \
500 (((seg_) = arena_seg_for_idx((arena_), __seg_idx)), true); \
501 __seg_idx++)
502
503void *arena_alloc_full(struct arena *a, size_t size, struct alloc_params ap)
504 cw_alloc(2) cc_warn_unused_result;
505
506void *arena_realloc_full(struct arena *a, void *ptr, size_t size,
507 struct alloc_params ap) cw_alloc(3)
508 cc_warn_unused_result;
509
510void *arena_alloc_special_full(struct arena *a, struct arena_params *p,
511 struct alloc_params ap) cw_alloc()
512 cc_warn_unused_result;
513
514void *arena_alloc_aligned_full(struct arena *a, size_t size, size_t align,
515 struct alloc_params ap) cw_alloc(2)
516 cc_warn_unused_result;
517
518err_checked arena_free_full(struct arena *a, void *ptr, struct alloc_params ap);
519
520err_checked arena_free_aligned_full(struct arena *a, void *ptr, size_t align,
521 struct alloc_params ap);
522
523err_checked arena_clear_full(struct arena *a, struct arena_params *p);
524
525void arena_destroy_full(struct arena *a, struct arena_params *p);
526
527void arena_reset_full(struct arena *a, struct arena_params *p);
528
529err_checked arena_set_tag_full(struct arena *a, void *ptr, arena_tag_t tag);
530
531err_checked arena_set_budget_full(struct arena *a, struct arena_budget b);
532
533size_t arena_qry_size_full(struct arena *a, void *p);
534
535arena_tag_t arena_qry_tag_full(struct arena *a, void *p);
536
537stack_handle_t arena_qry_trace_full(struct arena *a, void *p);
538
539bool arena_qry_owns_full(struct arena *a, const void *p);
540
541struct log_site *arena_get_log_site_full(struct arena *a,
542 struct arena_params *p);
543
544struct arena_budget arena_get_budget_full(struct arena *a,
545 struct arena_params *p);
546
547struct arena_dumpster *arena_get_dumpster_full(struct arena *a,
548 struct arena_params *p);
549
550arena_identity_t arena_get_identity_full(struct arena *a,
551 struct arena_params *p);
552
553bool arena_identity_eq_full(struct arena *a, struct arena *b);
554
555err_checked arena_on_hook_full(struct arena *a, enum arena_hook_type hook);
556
557bool arena_is_empty_full(struct arena *a);
558
559bool arena_getref_full(struct arena *a);
560
561void arena_putref_full(struct arena *a);
562
563struct arena *arena_create_via(struct arena_desc *desc, struct arena_params *p);
564
565/* TODO: Figure out dying semantics */
566static inline bool arena_bin_getref_rcu(struct arena_bin *bin) {
567 return refcount_inc_not_zero(rc: &bin->refcount);
568}
569
570static inline void arena_bin_putref(struct arena_bin *bin) {
571 cc_ignore(bin);
572}
573