1 /* SPDX-License-Identifier: (GPL-2.0-only OR BSD-2-Clause) */ 2 /* Copyright (c) 2024 Meta Platforms, Inc. and affiliates. */ 3 #pragma once 4 #include <errno.h> 5 #include "bpf_arena_alloc.h" 6 #include "bpf_arena_list.h" 7 8 struct htab_bucket { 9 struct arena_list_head head; 10 }; 11 typedef struct htab_bucket __arena htab_bucket_t; 12 13 struct htab { 14 htab_bucket_t *buckets; 15 int n_buckets; 16 }; 17 typedef struct htab __arena htab_t; 18 19 static inline htab_bucket_t *__select_bucket(htab_t *htab, __u32 hash) 20 { 21 htab_bucket_t *b = htab->buckets; 22 23 cast_kern(b); 24 return &b[hash & (htab->n_buckets - 1)]; 25 } 26 27 static inline arena_list_head_t *select_bucket(htab_t *htab, __u32 hash) 28 { 29 return &__select_bucket(htab, hash)->head; 30 } 31 32 struct hashtab_elem { 33 int hash; 34 int key; 35 int value; 36 struct arena_list_node hash_node; 37 }; 38 typedef struct hashtab_elem __arena hashtab_elem_t; 39 40 static hashtab_elem_t *lookup_elem_raw(arena_list_head_t *head, __u32 hash, int key) 41 { 42 hashtab_elem_t *l; 43 44 list_for_each_entry(l, head, hash_node) 45 if (l->hash == hash && l->key == key) 46 return l; 47 48 return NULL; 49 } 50 51 static int htab_hash(int key) 52 { 53 return key; 54 } 55 56 __weak int htab_lookup_elem(htab_t *htab __arg_arena, int key) 57 { 58 hashtab_elem_t *l_old; 59 arena_list_head_t *head; 60 61 cast_kern(htab); 62 head = select_bucket(htab, key); 63 l_old = lookup_elem_raw(head, htab_hash(key), key); 64 if (l_old) 65 return l_old->value; 66 return 0; 67 } 68 69 __weak int htab_update_elem(htab_t *htab __arg_arena, int key, int value) 70 { 71 hashtab_elem_t *l_new = NULL, *l_old; 72 arena_list_head_t *head; 73 74 cast_kern(htab); 75 head = select_bucket(htab, key); 76 l_old = lookup_elem_raw(head, htab_hash(key), key); 77 78 l_new = bpf_alloc(sizeof(*l_new)); 79 if (!l_new) 80 return -ENOMEM; 81 l_new->key = key; 82 l_new->hash = htab_hash(key); 83 l_new->value = value; 84 85 list_add_head(&l_new->hash_node, head); 86 if (l_old) { 87 list_del(&l_old->hash_node); 88 bpf_free(l_old); 89 } 90 return 0; 91 } 92 93 void htab_init(htab_t *htab) 94 { 95 void __arena *buckets = bpf_arena_alloc_pages(&arena, NULL, 2, NUMA_NO_NODE, 0); 96 97 cast_user(buckets); 98 htab->buckets = buckets; 99 htab->n_buckets = 2 * PAGE_SIZE / sizeof(struct htab_bucket); 100 } 101