tools/testing/selftests/bpf/bpf_arena_htab.h

   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 }