101 lines · c
1/* SPDX-License-Identifier: (GPL-2.0-only OR BSD-2-Clause) */2/* Copyright (c) 2024 Meta Platforms, Inc. and affiliates. */3#pragma once4#include <errno.h>5#include "bpf_arena_alloc.h"6#include "bpf_arena_list.h"7 8struct htab_bucket {9 struct arena_list_head head;10};11typedef struct htab_bucket __arena htab_bucket_t;12 13struct htab {14 htab_bucket_t *buckets;15 int n_buckets;16};17typedef struct htab __arena htab_t;18 19static 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 27static inline arena_list_head_t *select_bucket(htab_t *htab, __u32 hash)28{29 return &__select_bucket(htab, hash)->head;30}31 32struct hashtab_elem {33 int hash;34 int key;35 int value;36 struct arena_list_node hash_node;37};38typedef struct hashtab_elem __arena hashtab_elem_t;39 40static 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 51static 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 93void 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