xref: /linux/tools/testing/selftests/bpf/bpf_arena_htab.h (revision 79790b6818e96c58fe2bffee1b418c16e64e7b80)
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 
__select_bucket(htab_t * htab,__u32 hash)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 
select_bucket(htab_t * htab,__u32 hash)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 
lookup_elem_raw(arena_list_head_t * head,__u32 hash,int key)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 
htab_hash(int key)51 static int htab_hash(int key)
52 {
53 	return key;
54 }
55 
htab_lookup_elem(htab_t * htab __arg_arena,int key)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 
htab_update_elem(htab_t * htab __arg_arena,int key,int value)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 
htab_init(htab_t * htab)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