1 /* SPDX-License-Identifier: (LGPL-2.1 OR BSD-2-Clause) */ 2 #ifndef __TASK_LOCAL_DATA_H 3 #define __TASK_LOCAL_DATA_H 4 5 #include <errno.h> 6 #include <fcntl.h> 7 #include <sched.h> 8 #include <stdatomic.h> 9 #include <stddef.h> 10 #include <stdlib.h> 11 #include <string.h> 12 #include <unistd.h> 13 #include <sys/syscall.h> 14 #include <sys/types.h> 15 16 #ifdef TLD_FREE_DATA_ON_THREAD_EXIT 17 #include <pthread.h> 18 #endif 19 20 #include <bpf/bpf.h> 21 22 /* 23 * OPTIONS 24 * 25 * Define the option before including the header. Using different options in 26 * different translation units is strongly discouraged. 27 * 28 * TLD_FREE_DATA_ON_THREAD_EXIT - Frees memory on thread exit automatically 29 * 30 * Thread-specific memory for storing TLD is allocated lazily on the first call to 31 * tld_get_data(). The thread that calls it must also call tld_free() on thread exit 32 * to prevent memory leak. Pthread will be included if the option is defined. A pthread 33 * key will be registered with a destructor that calls tld_free(). Enabled only when 34 * the option is defined and TLD_DEFINE_KEY/tld_create_key() is called in the same 35 * translation unit. 36 * 37 * 38 * TLD_DYN_DATA_SIZE - The maximum size of memory allocated for TLDs created dynamically 39 * (default: 64 bytes) 40 * 41 * A TLD can be defined statically using TLD_DEFINE_KEY() or created on the fly using 42 * tld_create_key(). As the total size of TLDs created with tld_create_key() cannot be 43 * possibly known statically, a memory area of size TLD_DYN_DATA_SIZE will be allocated 44 * for these TLDs. This additional memory is allocated for every thread that calls 45 * tld_get_data() even if no tld_create_key are actually called, so be mindful of 46 * potential memory wastage. Use TLD_DEFINE_KEY() whenever possible as just enough memory 47 * will be allocated for TLDs created with it. 48 * 49 * 50 * TLD_NAME_LEN - The maximum length of the name of a TLD (default: 62) 51 * 52 * Setting TLD_NAME_LEN will affect the maximum number of TLDs a process can store, 53 * TLD_MAX_DATA_CNT. Must be consistent with task_local_data.bpf.h. 54 * 55 * 56 * TLD_DONT_ROUND_UP_DATA_SIZE - Don't round up memory size allocated for data if 57 * the memory allocator has low overhead aligned_alloc() implementation. 58 * 59 * For some memory allocators, when calling aligned_alloc(alignment, size), size 60 * does not need to be an integral multiple of alignment and it can be fulfilled 61 * without using round_up(size, alignment) bytes of memory. Enable this option to 62 * reduce memory usage. 63 */ 64 65 #define TLD_PAGE_SIZE getpagesize() 66 #define TLD_PAGE_MASK (~(TLD_PAGE_SIZE - 1)) 67 68 #define TLD_ROUND_MASK(x, y) ((__typeof__(x))((y) - 1)) 69 #define TLD_ROUND_UP(x, y) ((((x) - 1) | TLD_ROUND_MASK(x, y)) + 1) 70 71 #define TLD_ROUND_UP_POWER_OF_TWO(x) (1UL << (sizeof(x) * 8 - __builtin_clzl(x - 1))) 72 73 #ifndef TLD_DYN_DATA_SIZE 74 #define TLD_DYN_DATA_SIZE 64 75 #endif 76 77 #define TLD_MAX_DATA_CNT (TLD_PAGE_SIZE / sizeof(struct tld_metadata) - 1) 78 79 #ifndef TLD_NAME_LEN 80 #define TLD_NAME_LEN 62 81 #endif 82 83 #ifdef __cplusplus 84 extern "C" { 85 #endif 86 87 typedef struct { 88 __s16 off; 89 } tld_key_t; 90 91 struct tld_metadata { 92 char name[TLD_NAME_LEN]; 93 _Atomic __u16 size; /* size of tld_data_u->data */ 94 }; 95 96 struct tld_meta_u { 97 _Atomic __u16 cnt; 98 __u16 size; 99 struct tld_metadata metadata[]; 100 }; 101 102 struct tld_data_u { 103 __u64 start; /* offset of tld_data_u->data in a page */ 104 char data[] __attribute__((aligned(8))); 105 }; 106 107 struct tld_map_value { 108 void *data; 109 struct tld_meta_u *meta; 110 }; 111 112 struct tld_meta_u * _Atomic tld_meta_p __attribute__((weak)); 113 __thread struct tld_data_u *tld_data_p __attribute__((weak)); 114 115 #ifdef TLD_FREE_DATA_ON_THREAD_EXIT 116 bool _Atomic tld_pthread_key_init __attribute__((weak)); 117 pthread_key_t tld_pthread_key __attribute__((weak)); 118 119 static void tld_free(void); 120 121 static void __tld_thread_exit_handler(void *unused) 122 { 123 (void)unused; 124 tld_free(); 125 } 126 #endif 127 128 static int __tld_init_meta_p(void) 129 { 130 struct tld_meta_u *meta, *uninit = NULL; 131 int err = 0; 132 133 meta = (struct tld_meta_u *)aligned_alloc(TLD_PAGE_SIZE, TLD_PAGE_SIZE); 134 if (!meta) { 135 err = -ENOMEM; 136 goto out; 137 } 138 139 memset(meta, 0, TLD_PAGE_SIZE); 140 meta->size = TLD_DYN_DATA_SIZE; 141 142 if (!atomic_compare_exchange_strong(&tld_meta_p, &uninit, meta)) { 143 free(meta); 144 goto out; 145 } 146 147 out: 148 return err; 149 } 150 151 static int __tld_init_data_p(int map_fd) 152 { 153 struct tld_map_value map_val; 154 struct tld_data_u *data; 155 int err, tid_fd = -1; 156 size_t size, size_pot; 157 158 tid_fd = syscall(SYS_pidfd_open, sys_gettid(), O_EXCL); 159 if (tid_fd < 0) { 160 err = -errno; 161 goto out; 162 } 163 164 /* 165 * tld_meta_p->size = TLD_DYN_DATA_SIZE + 166 * total size of TLDs defined via TLD_DEFINE_KEY() 167 */ 168 size = tld_meta_p->size + sizeof(struct tld_data_u); 169 size_pot = TLD_ROUND_UP_POWER_OF_TWO(size); 170 #ifdef TLD_DONT_ROUND_UP_DATA_SIZE 171 data = (struct tld_data_u *)aligned_alloc(size_pot, size); 172 #else 173 data = (struct tld_data_u *)aligned_alloc(size_pot, size_pot); 174 #endif 175 if (!data) { 176 err = -ENOMEM; 177 goto out; 178 } 179 180 /* 181 * Always pass a page-aligned address to UPTR since the size of tld_map_value::data 182 * is a page in BTF. 183 */ 184 map_val.data = (void *)(TLD_PAGE_MASK & (intptr_t)data); 185 data->start = (~TLD_PAGE_MASK & (intptr_t)data) + sizeof(struct tld_data_u); 186 map_val.meta = tld_meta_p; 187 188 err = bpf_map_update_elem(map_fd, &tid_fd, &map_val, 0); 189 if (err) { 190 free(data); 191 goto out; 192 } 193 194 tld_data_p = data; 195 #ifdef TLD_FREE_DATA_ON_THREAD_EXIT 196 pthread_setspecific(tld_pthread_key, (void *)1); 197 #endif 198 out: 199 if (tid_fd >= 0) 200 close(tid_fd); 201 return err; 202 } 203 204 static tld_key_t __tld_create_key(const char *name, size_t size, bool dyn_data) 205 { 206 int err, i, sz, off = 0; 207 bool uninit = false; 208 __u16 cnt; 209 210 if (!tld_meta_p) { 211 err = __tld_init_meta_p(); 212 if (err) 213 return (tld_key_t){(__s16)err}; 214 } 215 216 #ifdef TLD_FREE_DATA_ON_THREAD_EXIT 217 if (atomic_compare_exchange_strong(&tld_pthread_key_init, &uninit, true)) { 218 err = pthread_key_create(&tld_pthread_key, __tld_thread_exit_handler); 219 if (err) 220 return (tld_key_t){(__s16)err}; 221 } 222 #endif 223 224 for (i = 0; i < (int)TLD_MAX_DATA_CNT; i++) { 225 retry: 226 cnt = atomic_load(&tld_meta_p->cnt); 227 if (i < cnt) { 228 /* A metadata is not ready until size is updated with a non-zero value */ 229 while (!(sz = atomic_load(&tld_meta_p->metadata[i].size))) 230 sched_yield(); 231 232 if (!strncmp(tld_meta_p->metadata[i].name, name, TLD_NAME_LEN)) 233 return (tld_key_t){-EEXIST}; 234 235 off += TLD_ROUND_UP(sz, 8); 236 continue; 237 } 238 239 /* 240 * TLD_DEFINE_KEY() is given memory upto a page while at most 241 * TLD_DYN_DATA_SIZE is allocated for tld_create_key() 242 */ 243 if (dyn_data) { 244 if (off + TLD_ROUND_UP(size, 8) > tld_meta_p->size) 245 return (tld_key_t){-E2BIG}; 246 } else { 247 if (off + TLD_ROUND_UP(size, 8) > TLD_PAGE_SIZE - sizeof(struct tld_data_u)) 248 return (tld_key_t){-E2BIG}; 249 tld_meta_p->size += TLD_ROUND_UP(size, 8); 250 } 251 252 /* 253 * Only one tld_create_key() can increase the current cnt by one and 254 * takes the latest available slot. Other threads will check again if a new 255 * TLD can still be added, and then compete for the new slot after the 256 * succeeding thread update the size. 257 */ 258 if (!atomic_compare_exchange_strong(&tld_meta_p->cnt, &cnt, cnt + 1)) 259 goto retry; 260 261 strscpy(tld_meta_p->metadata[i].name, name); 262 atomic_store(&tld_meta_p->metadata[i].size, size); 263 return (tld_key_t){(__s16)off}; 264 } 265 266 return (tld_key_t){-ENOSPC}; 267 } 268 269 /** 270 * TLD_DEFINE_KEY() - Define a TLD and a global variable key associated with the TLD. 271 * 272 * @name: The name of the TLD 273 * @size: The size of the TLD 274 * @key: The variable name of the key. Cannot exceed TLD_NAME_LEN 275 * 276 * The macro can only be used in file scope. 277 * 278 * A global variable key of opaque type, tld_key_t, will be declared and initialized before 279 * main() starts. Use tld_key_is_err() or tld_key_err_or_zero() later to check if the key 280 * creation succeeded. Pass the key to tld_get_data() to get a pointer to the TLD. 281 * bpf programs can also fetch the same key by name. 282 * 283 * The total size of TLDs created using TLD_DEFINE_KEY() cannot exceed a page. Just 284 * enough memory will be allocated for each thread on the first call to tld_get_data(). 285 */ 286 #define TLD_DEFINE_KEY(key, name, size) \ 287 tld_key_t key; \ 288 \ 289 __attribute__((constructor(101))) \ 290 void __tld_define_key_##key(void) \ 291 { \ 292 key = __tld_create_key(name, size, false); \ 293 } 294 295 /** 296 * tld_create_key() - Create a TLD and return a key associated with the TLD. 297 * 298 * @name: The name the TLD 299 * @size: The size of the TLD 300 * 301 * Return an opaque object key. Use tld_key_is_err() or tld_key_err_or_zero() to check 302 * if the key creation succeeded. Pass the key to tld_get_data() to get a pointer to 303 * locate the TLD. bpf programs can also fetch the same key by name. 304 * 305 * Use tld_create_key() only when a TLD needs to be created dynamically (e.g., @name is 306 * not known statically or a TLD needs to be created conditionally) 307 * 308 * An additional TLD_DYN_DATA_SIZE bytes are allocated per-thread to accommodate TLDs 309 * created dynamically with tld_create_key(). Since only a user page is pinned to the 310 * kernel, when TLDs created with TLD_DEFINE_KEY() uses more than TLD_PAGE_SIZE - 311 * TLD_DYN_DATA_SIZE, the buffer size will be limited to the rest of the page. 312 */ 313 __attribute__((unused)) 314 static tld_key_t tld_create_key(const char *name, size_t size) 315 { 316 return __tld_create_key(name, size, true); 317 } 318 319 __attribute__((unused)) 320 static inline bool tld_key_is_err(tld_key_t key) 321 { 322 return key.off < 0; 323 } 324 325 __attribute__((unused)) 326 static inline int tld_key_err_or_zero(tld_key_t key) 327 { 328 return tld_key_is_err(key) ? key.off : 0; 329 } 330 331 /** 332 * tld_get_data() - Get a pointer to the TLD associated with the given key of the 333 * calling thread. 334 * 335 * @map_fd: A file descriptor of tld_data_map, the underlying BPF task local storage map 336 * of task local data. 337 * @key: A key object created by TLD_DEFINE_KEY() or tld_create_key(). 338 * 339 * Return a pointer to the TLD if the key is valid; NULL if not enough memory for TLD 340 * for this thread, or the key is invalid. The returned pointer is guaranteed to be 8-byte 341 * aligned. 342 * 343 * Threads that call tld_get_data() must call tld_free() on exit to prevent 344 * memory leak if TLD_FREE_DATA_ON_THREAD_EXIT is not defined. 345 */ 346 __attribute__((unused)) 347 static void *tld_get_data(int map_fd, tld_key_t key) 348 { 349 if (!tld_meta_p) 350 return NULL; 351 352 /* tld_data_p is allocated on the first invocation of tld_get_data() */ 353 if (!tld_data_p && __tld_init_data_p(map_fd)) 354 return NULL; 355 356 return tld_data_p->data + key.off; 357 } 358 359 /** 360 * tld_free() - Free task local data memory of the calling thread 361 * 362 * For the calling thread, all pointers to TLDs acquired before will become invalid. 363 * 364 * Users must call tld_free() on thread exit to prevent memory leak. Alternatively, 365 * define TLD_FREE_DATA_ON_THREAD_EXIT and a thread exit handler will be registered 366 * to free the memory automatically. Calling tld_free() before thread exit is 367 * undefined behavior, which may lead to null-pointer dereference. 368 */ 369 __attribute__((unused)) 370 static void tld_free(void) 371 { 372 if (tld_data_p) { 373 free(tld_data_p); 374 tld_data_p = NULL; 375 } 376 } 377 378 #ifdef __cplusplus 379 } /* extern "C" */ 380 #endif 381 382 #endif /* __TASK_LOCAL_DATA_H */ 383