1#include <math/fixed.h>
2#include <math/hash.h>
3#include <math/sort.h>
4#include <mem/alloc.h>
5#include <mem/arena_provider.h>
6#include <structures/list_rcu.h>
7#include <sync/rcu.h>
8
9#include "internal.h"
10
11struct arena_globals arena_global = {0};
12enum err arena_budget_prio_scale(sz_b_t used, int scale,
13 ALLOC_PRIORITY_BITMAP_DECLARE(prio_map_out)) {
14 fx32_32_t y = fx_pow_i32(FX_E, exp: (int) used + scale);
15 if (y > fx_from_int(x: ALLOC_PRIORITY_MAX) || y < FX_ZERO)
16 return ERR_OVERFLOW;
17
18 bitmap_zero(map: prio_map_out, nbits: ALLOC_PRIORITY_MAX);
19 uint16_t base = (uint16_t) fx_to_int(x: y);
20 for (uint16_t i = base; i < ALLOC_PRIORITY_MAX; i++)
21 bitmap_set(map: prio_map_out, bit: i);
22
23 return ERR_OK;
24}
25
26static void arena_desc_table_free_cb(struct rcu_cb *cb) {
27 struct arena_desc_table *tbl =
28 container_of(cb, struct arena_desc_table, rcu_cb);
29 vfree(tbl->descs);
30 vfree(tbl);
31}
32
33static int arena_desc_cmp(const void *a, const void *b) {
34 const struct arena_desc *ada = *(struct arena_desc *const *) a;
35 const struct arena_desc *adb = *(struct arena_desc *const *) b;
36 return (ada->strategy > adb->strategy) - (ada->strategy < adb->strategy);
37}
38
39static int arena_desc_strategy_cmp(const void *key, const void *elem) {
40 enum arena_strategy strat = *(const enum arena_strategy *) key;
41 const struct arena_desc *d = *(struct arena_desc *const *) elem;
42 return (strat > d->strategy) - (strat < d->strategy);
43}
44
45enum err arena_desc_register(struct arena_desc *desc) {
46 enum err err = ERR_OK;
47
48 enum irql irql = spin_lock(&arena_global.tbl_lock);
49
50 struct arena_desc_table *old_tbl = rcu_dereference(arena_global.desc_table);
51 enforce(old_tbl);
52
53 size_t old_desc_count = old_tbl->desc_count;
54 size_t desc_count = old_desc_count + 1;
55 struct arena_desc_table *new_tbl = vnew(new_tbl, ALLOC_ZERO);
56 if (!new_tbl) {
57 err = ERR_NO_MEM;
58 goto out;
59 }
60
61 struct arena_desc **new_descs = vnews(new_descs, desc_count, ALLOC_ZERO);
62 if (!new_descs) {
63 err = ERR_NO_MEM;
64 vfree(new_tbl);
65 goto out;
66 }
67
68 if (old_desc_count && old_tbl->descs)
69 memcpy(new_descs, old_tbl->descs,
70 old_desc_count * sizeof(struct arena_desc *));
71
72 new_descs[old_desc_count] = desc;
73 heapsort(vbase: new_descs, nmemb: desc_count, size: sizeof(*new_descs), compar: arena_desc_cmp);
74 new_tbl->desc_count = desc_count;
75 new_tbl->descs = new_descs;
76
77 rcu_assign_pointer(arena_global.desc_table, new_tbl);
78 rcu_defer(cb: &old_tbl->rcu_cb, func: arena_desc_table_free_cb);
79
80out:
81 spin_unlock(&arena_global.tbl_lock, irql);
82 return err;
83}
84
85struct arena_desc *arena_desc_lookup(enum arena_strategy strat) {
86 rcu_read_lock();
87
88 struct arena_desc_table *tbl = rcu_dereference(arena_global.desc_table);
89
90 struct arena_desc **slot =
91 tbl->descs ? bsearch(key: &strat, base: tbl->descs, nmemb: tbl->desc_count,
92 size: sizeof(*tbl->descs), compar: arena_desc_strategy_cmp)
93 : NULL;
94 struct arena_desc *found = slot ? *slot : NULL;
95
96 rcu_read_unlock();
97
98 return found;
99}
100
101void arena_global_init(void) {
102 spinlock_init(&arena_global.tbl_lock);
103 struct arena_desc_table *tbl = vnew(tbl, ALLOC_ZERO);
104
105 rcu_assign_pointer(arena_global.desc_table, tbl);
106}
107
108static uint16_t size_for_id(struct arena_seg_desc *seg_descs, size_t seg_count,
109 uint16_t id) {
110 for (size_t i = 0; i < seg_count; i++) {
111 if (seg_descs[i].id == id)
112 return seg_descs[i].size;
113 }
114
115 unreachable("invalid id");
116}
117
118/* Errors panic because the provider calls this */
119struct arena *arena_create_full(struct arena_seg_desc *seg_descs,
120 size_t seg_count) {
121 size_t seg_data_size_total = 0;
122 struct arena_seg_desc *large_desc = NULL;
123 BITMAP_DECLARE(seen_id_bitmap, ARENA_MAX_SEG) = {0};
124 for (size_t i = 0; i < seg_count; i++) {
125 enforce(seg_descs[i].size);
126
127 /* It's fine if a large segment is smaller than ARENA_SEG_MAX_SIZE */
128 if (seg_descs[i].large) {
129 enforce(!large_desc, "more than one large segment exists");
130 large_desc = &seg_descs[i];
131 }
132
133 if (seg_descs[i].size >= ARENA_SEG_MAX_SIZE)
134 enforce(seg_descs[i].large, "large segment id %u must set .large",
135 seg_descs[i].id);
136
137 enforce(!bitmap_test_and_set(seen_id_bitmap, seg_descs[i].id),
138 DUPLICATE, "id %hu", seg_descs[i].id);
139
140 seg_data_size_total += ALIGN_UP(seg_descs[i].size, ARENA_SEG_ALIGN);
141 }
142
143 size_t inmem_size_raw = seg_count * sizeof(struct arena_seg_inmem_desc);
144 size_t inmem_size = ALIGN_UP(inmem_size_raw, ARENA_SEG_ALIGN);
145 size_t seg_size_total = inmem_size + seg_data_size_total;
146 size_t size_total = sizeof(struct arena) + seg_size_total;
147
148 struct arena *arena = valloc(size_total, ALLOC_ZERO);
149 if (!arena)
150 return NULL;
151
152#ifdef DEBUG_ARENA
153 arena->seg_count = seg_count;
154 arena->total_size = size_total;
155#endif
156
157 struct arena_seg_inmem_desc *inmem_descs = arena_get_inmem_descs(a: arena);
158
159 if (large_desc) {
160 inmem_descs[seg_count - 1].id = large_desc->id;
161
162#ifdef DEBUG_ARENA
163 inmem_descs[seg_count - 1].present = true;
164 inmem_descs[seg_count - 1].type = large_desc->type;
165 inmem_descs[seg_count - 1].size = large_desc->size;
166#endif
167 }
168
169 size_t cursor = 0;
170 for (size_t i = 0; i < seg_count; i++) {
171 if (!seg_descs[i].large) {
172
173#ifdef DEBUG_ARENA
174 inmem_descs[cursor].present = true;
175 inmem_descs[cursor].size = seg_descs[i].size;
176 inmem_descs[cursor].type = seg_descs[i].type;
177#endif
178
179 inmem_descs[cursor].id = seg_descs[i].id;
180 cursor++;
181 }
182 }
183
184 size_t cc_unused cursor_offset = 0;
185 for (uint16_t i = 0; i < seg_count; i++) {
186 uint16_t size;
187
188 if (i == 0) {
189 size = inmem_size;
190 } else {
191 uint16_t id = inmem_descs[i - 1].id;
192 size = ALIGN_UP(size_for_id(seg_descs, seg_count, id),
193 ARENA_SEG_ALIGN);
194 }
195
196 cursor_offset += size;
197
198#ifdef DEBUG_ARENA
199 inmem_descs[i].seg =
200 (struct arena_seg *) &arena->payload[cursor_offset];
201#endif
202
203 inmem_descs[i].offset_bump = size;
204 }
205
206 arena_check_assert(arena);
207 return arena;
208}
209
210size_t arena_seg_count(struct arena *a) {
211 struct arena_seg_inmem_desc *descs = arena_get_inmem_descs(a);
212
213#ifdef DEBUG_ARENA
214 enforce(a->seg_count,
215 "attempted to count segments of arena with no segments");
216#endif
217
218 size_t count = 0;
219 uint16_t end = descs[0].offset_bump / sizeof(struct arena_seg_inmem_desc);
220 for (uint16_t i = 0; i < end; i++) {
221 struct arena_seg_inmem_desc *desc = &descs[i];
222 if (desc->offset_bump)
223 count++;
224 }
225
226 return count;
227}
228
229struct arena_seg arena_seg_lookup(struct arena *a, uint16_t seg_id) {
230 struct arena_seg_inmem_desc *inmem_descs = arena_get_inmem_descs(a);
231
232#ifdef DEBUG_ARENA
233 bool found = false;
234 for (uint16_t i = 0; i < a->seg_count; i++) {
235 if (inmem_descs[i].id == seg_id)
236 found = true;
237 }
238
239 enforce(found, NOT_FOUND, "segment %u", seg_id);
240#endif
241
242 size_t cc_unused cursor = 0;
243 for (int i = 0; i < ARENA_MAX_SEG; i++) {
244 cursor += inmem_descs[i].offset_bump;
245 if (inmem_descs[i].id == seg_id) {
246 uint8_t *storage = &a->payload[cursor];
247 return (struct arena_seg){.storage = storage,
248 .id = inmem_descs[i].id};
249 }
250 }
251
252 panic(CORRUPTION, "segment %u vanished during traversal", seg_id);
253}
254
255struct arena_seg arena_seg_for_idx(struct arena *a, uint16_t idx) {
256#ifdef DEBUG_ARENA
257 enforce(idx < a->seg_count, "index %u out of bounds of %u", idx,
258 a->seg_count);
259#endif
260
261 struct arena_seg_inmem_desc *inmem_descs = arena_get_inmem_descs(a);
262
263 size_t cursor = 0;
264 for (int i = 0; i < ARENA_MAX_SEG; i++) {
265 cursor += inmem_descs[i].offset_bump;
266 if (i == idx) {
267 uint8_t *storage = &a->payload[cursor];
268 return (struct arena_seg){.storage = storage,
269 .id = inmem_descs[i].id};
270 }
271 }
272
273 unreachable();
274}
275
276static struct arena_bucket *get_bucket(struct arena_dumpster *ad,
277 arena_identity_t id) {
278 uint32_t hash = hash_jenkins_qword(key: id, /* TODO: seed = */ seed: 12345);
279 size_t idx = hash % ad->bucket_count;
280
281 struct arena_bucket *bkt = &ad->buckets[idx];
282 return enforce(bkt);
283}
284
285static struct arena_bin *bucket_id_lookup(struct arena_bucket *bkt,
286 arena_identity_t id) {
287 struct arena_bin *entry, *leader = NULL;
288 hlist_for_each_entry_rcu(entry, &bkt->bins, hlist_node) {
289 if (entry->identity == id) {
290 leader = entry;
291 break;
292 }
293 }
294
295 return leader;
296}
297
298/*
299 * Memory arena buckets have physical and acting leaders. For every bucket,
300 * we have an hlist that looks something like
301 *
302 * bucket
303 * |
304 * +-> hlist_head -> node -> node -> node -> NULL
305 *
306 * The nodes on this list are the "physical" leaders. However, the "physical"
307 * leaders are merely an optimization, because the logical leaders are the
308 * "acting" leaders.
309 *
310 * The reason we have this split at all is because any traversal of a bucket
311 * leader list necessarily races with simultaneous deletion of arena bins,
312 * i.e. marking the arena bin as dying and dropping its reference count.
313 *
314 * Such bins are NOT to be used any more, unless someone owned a reference
315 * count prior to the start of the destruction.
316 *
317 * Thus, any traversal treats the leader as being the "first bin in the list
318 * of the physical leader that we could successfully acquire a reference to".
319 *
320 * Thus, if a physical leader is alive, it is also the acting leader.
321 * However, if a physical leader is dying/has died, but is still momentarily
322 * visible on the list in the read-side critical section, as the RCU
323 * free callback has not run, we know that the next-in-line is the acting
324 * leader, which is where the arenas would get freed to in a dumpster.
325 */
326static struct arena_bin *arena_dumpster_find_leader(struct arena_dumpster *ad,
327 arena_identity_t id) {
328 struct arena_bucket *bkt = get_bucket(ad, id);
329 struct arena_bin *entry, **found = NULL;
330
331 rcu_read_lock();
332
333 struct arena_bin *first_leader = bucket_id_lookup(bkt, id);
334
335 if (!first_leader)
336 goto out;
337
338 found = &first_leader;
339
340 /* list_del_rcu prevents the middle links from getting
341 * their own references wiped, so traversal under the
342 * leader's list is fine */
343 if (!arena_bin_getref_rcu(bin: first_leader)) {
344 found = NULL;
345 list_for_each_entry_rcu(
346 entry, list_node_as_head(&first_leader->bin_list), bin_list) {
347 if (arena_bin_getref_rcu(bin: entry)) {
348 found = &entry;
349 goto out;
350 }
351 }
352 }
353
354out:
355 rcu_read_unlock();
356
357 return found ? *found : NULL;
358}
359
360void arena_dumpster_add_bin(struct arena_dumpster *ad, struct arena_bin *ab) {
361 struct arena_bucket *bkt = get_bucket(ad, id: ab->identity);
362 struct mutex *mtx = &bkt->mutex;
363
364 mutex_lock(mtx);
365
366 struct arena_bin *leader = arena_dumpster_find_leader(ad, id: ab->identity);
367
368 if (!leader) {
369 list_node_init(n: &ab->bin_list);
370 hlist_add_head_rcu(n: &ab->hlist_node, h: &bkt->bins);
371 } else {
372 list_head_push_back_rcu(h: list_node_as_head(n: &leader->bin_list),
373 n: &ab->bin_list);
374 ab->not_head_magic = ARENA_BIN_NOT_HEAD_MAGIC;
375 }
376
377 ab->dumpster = ad;
378 mutex_unlock(mtx);
379
380 if (leader)
381 arena_bin_putref(bin: leader);
382}
383
384void arena_dumpster_delete_bin(struct arena_dumpster *ad,
385 struct arena_bin *ab) {
386 struct arena_bucket *bkt = get_bucket(ad, id: ab->identity);
387 struct mutex *mtx = &bkt->mutex;
388
389 mutex_guard(mtx);
390
391 struct arena_bin *leader = bucket_id_lookup(bkt, id: ab->identity);
392 enforce(leader, BAD_FREE, "bin has no leader");
393
394 list_node_delete_rcu(n: &ab->bin_list);
395 if (leader == ab) {
396 struct arena_bin *iter, *next = NULL;
397 list_for_each_entry_rcu(iter, list_node_as_head(&leader->bin_list),
398 bin_list) {
399 /* This is the next physical leader */
400 if (arena_bin_getref_rcu(bin: iter)) {
401 next = iter;
402 break;
403 }
404 }
405
406 if (next)
407 hlist_add_head_rcu(n: &next->hlist_node, h: &bkt->bins);
408 hlist_del_rcu(n: &ab->hlist_node);
409 }
410}
411