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 18 static inline htab_bucket_t *__select_bucket(struct htab __arena *htab, __u32 hash) 19 { 20 htab_bucket_t *b = htab->buckets; 21 22 cast_kern(b); 23 return &b[hash & (htab->n_buckets - 1)]; 24 } 25 26 static inline arena_list_head_t *select_bucket(struct htab __arena *htab, __u32 hash) 27 { 28 return &__select_bucket(htab, hash)->head; 29 } 30 31 struct hashtab_elem { 32 int hash; 33 int key; 34 int value; 35 struct arena_list_node hash_node; 36 }; 37 typedef struct hashtab_elem __arena hashtab_elem_t; 38 39 static hashtab_elem_t *lookup_elem_raw(arena_list_head_t *head, __u32 hash, int key) 40 { 41 hashtab_elem_t *l; 42 43 list_for_each_entry(l, head, hash_node) 44 if (l->hash == hash && l->key == key) 45 return l; 46 47 return NULL; 48 } 49 50 static int htab_hash(int key) 51 { 52 return key; 53 } 54 55 __weak int htab_lookup_elem(struct htab __arena *htab, int key) 56 { 57 hashtab_elem_t *l_old; 58 arena_list_head_t *head; 59 60 cast_kern(htab); 61 head = select_bucket(htab, key); 62 l_old = lookup_elem_raw(head, htab_hash(key), key); 63 if (l_old) 64 return l_old->value; 65 return 0; 66 } 67 68 __weak int htab_update_elem(struct htab __arena *htab, int key, int value) 69 { 70 hashtab_elem_t *l_new = NULL, *l_old; 71 arena_list_head_t *head; 72 73 cast_kern(htab); 74 head = select_bucket(htab, key); 75 l_old = lookup_elem_raw(head, htab_hash(key), key); 76 77 l_new = bpf_alloc(sizeof(*l_new)); 78 if (!l_new) 79 return -ENOMEM; 80 l_new->key = key; 81 l_new->hash = htab_hash(key); 82 l_new->value = value; 83 84 list_add_head(&l_new->hash_node, head); 85 if (l_old) { 86 list_del(&l_old->hash_node); 87 bpf_free(l_old); 88 } 89 return 0; 90 } 91 92 void htab_init(struct htab __arena *htab) 93 { 94 void __arena *buckets = bpf_arena_alloc_pages(&arena, NULL, 2, NUMA_NO_NODE, 0); 95 96 cast_user(buckets); 97 htab->buckets = buckets; 98 htab->n_buckets = 2 * PAGE_SIZE / sizeof(struct htab_bucket); 99 } 100