1 // SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause 2 /* 3 * Copyright (c) 2025-2026 Meta Platforms, Inc. and affiliates. 4 * Copyright (c) 2025-2026 Emil Tsalapatis <etsal@meta.com> 5 */ 6 7 #include <bpf_atomic.h> 8 9 #include <libarena/common.h> 10 11 #include <libarena/asan.h> 12 #include <libarena/spmc.h> 13 14 static inline 15 u64 spmc_arr_size(volatile struct spmc_arr __arena *spmc_arr) 16 { 17 return SPMC_ARR_BASESZ << spmc_arr->order; 18 } 19 20 static inline 21 u64 spmc_arr_get(volatile struct spmc_arr __arena *spmc_arr, u64 ind) 22 { 23 u64 ret = READ_ONCE(spmc_arr->data[ind % spmc_arr_size(spmc_arr)]); 24 25 return ret; 26 } 27 28 static inline 29 void spmc_arr_put(volatile struct spmc_arr __arena *spmc_arr, u64 ind, u64 value) 30 { 31 WRITE_ONCE(spmc_arr->data[ind % spmc_arr_size(spmc_arr)], value); 32 } 33 34 static inline 35 void spmc_arr_copy(volatile struct spmc_arr __arena *dst, 36 volatile struct spmc_arr __arena *src, u64 b, u64 t) 37 { 38 u64 i; 39 40 for (i = t; i < b && can_loop; i++) 41 spmc_arr_put(dst, i, spmc_arr_get(src, i)); 42 } 43 44 static inline 45 int spmc_order_init(struct spmc __arena *spmc, int order) 46 { 47 volatile struct spmc_arr __arena *arr = &spmc->arr[order]; 48 49 if (unlikely(!spmc)) 50 return -EINVAL; 51 52 if (order >= SPMC_ARR_ORDERS) 53 return -E2BIG; 54 55 /* Already allocated? */ 56 if (arr->data) 57 return 0; 58 59 arr->data = arena_malloc((SPMC_ARR_BASESZ << order) * sizeof(*arr->data)); 60 if (!arr->data) 61 return -ENOMEM; 62 63 return 0; 64 } 65 66 __weak 67 int spmc_owned_add(struct spmc __arena *spmc, u64 val) 68 { 69 volatile struct spmc_arr __arena *newarr; 70 volatile struct spmc_arr __arena *arr; 71 ssize_t sz; 72 u64 b, t; 73 int ret; 74 75 if (unlikely(!spmc)) 76 return -EINVAL; 77 78 /* 79 * Bottom must always be read first, also 80 * see spmc_steal(). 81 */ 82 b = smp_load_acquire(&spmc->bottom); 83 t = READ_ONCE(spmc->top); 84 arr = READ_ONCE(spmc->cur); 85 86 sz = b - t; 87 if (sz >= spmc_arr_size(arr) - 1) { 88 ret = spmc_order_init(spmc, arr->order + 1); 89 if (ret) 90 return ret; 91 92 newarr = &spmc->arr[arr->order + 1]; 93 94 spmc_arr_copy(newarr, arr, b, t); 95 smp_store_release(&spmc->cur, newarr); 96 arr = newarr; 97 } 98 99 spmc_arr_put(arr, b, val); 100 smp_store_release(&spmc->bottom, b + 1); 101 102 return 0; 103 } 104 105 106 __weak 107 int spmc_owned_remove(struct spmc __arena *spmc, u64 *val) 108 { 109 volatile struct spmc_arr __arena *arr; 110 int ret = 0; 111 ssize_t sz; 112 u64 value; 113 u64 b, t; 114 115 if (unlikely(!spmc || !val)) 116 return -EINVAL; 117 118 b = READ_ONCE(spmc->bottom) - 1; 119 WRITE_ONCE(spmc->bottom, b); 120 smp_mb(); 121 122 t = READ_ONCE(spmc->top); 123 arr = READ_ONCE(spmc->cur); 124 125 sz = b - t; 126 if (sz < 0) { 127 WRITE_ONCE(spmc->bottom, t); 128 return -ENOENT; 129 } 130 131 value = spmc_arr_get(arr, b); 132 if (sz > 0) { 133 *val = value; 134 return 0; 135 } 136 137 if (cmpxchg(&spmc->top, t, t + 1) != t) 138 ret = -EAGAIN; 139 140 WRITE_ONCE(spmc->bottom, t + 1); 141 142 if (ret) 143 return ret; 144 145 *val = value; 146 147 return 0; 148 } 149 150 __weak 151 int spmc_steal(struct spmc __arena *spmc, u64 *val) 152 { 153 volatile struct spmc_arr __arena *arr; 154 ssize_t sz; 155 u64 value; 156 u64 b, t; 157 158 if (unlikely(!spmc || !val)) 159 return -EINVAL; 160 161 /* 162 * It is important that t is read before b for 163 * stealers to avoid racing with the owner. 164 * Races between stealers are dealt with using 165 * CAS to increment the top value below. 166 */ 167 t = smp_load_acquire(&spmc->top); 168 b = smp_load_acquire(&spmc->bottom); 169 170 sz = b - t; 171 if (sz <= 0) 172 return -ENOENT; 173 174 arr = smp_load_acquire(&spmc->cur); 175 value = spmc_arr_get(arr, t); 176 177 if (cmpxchg(&spmc->top, t, t + 1) != t) 178 return -EAGAIN; 179 180 *val = value; 181 182 return 0; 183 } 184 185 186 __weak 187 struct spmc __arena *spmc_create(void) 188 { 189 /* 190 * Marked as volatile because otherwise the array 191 * reference in the internal loop gets demoted to 192 * scalar and the program fails verification. 193 */ 194 struct spmc __arena *volatile spmc; 195 int ret, i; 196 197 spmc = arena_malloc(sizeof(*spmc)); 198 if (!spmc) 199 return NULL; 200 201 spmc->bottom = 0; 202 spmc->top = 0; 203 204 for (i = 0; i < SPMC_ARR_ORDERS && can_loop; i++) { 205 spmc->arr[i].data = NULL; 206 spmc->arr[i].order = i; 207 } 208 209 ret = spmc_order_init((struct spmc __arena *)spmc, 0); 210 if (ret) { 211 arena_free(spmc); 212 return NULL; 213 } 214 215 spmc->cur = &spmc->arr[0]; 216 217 return (struct spmc __arena *)spmc; 218 } 219 220 __weak 221 int spmc_destroy(struct spmc __arena *spmc) 222 { 223 int i; 224 225 if (unlikely(!spmc)) 226 return -EINVAL; 227 228 for (i = 0; i < SPMC_ARR_ORDERS && can_loop; i++) 229 arena_free(spmc->arr[i].data); 230 231 arena_free(spmc); 232 233 return 0; 234 } 235