tools/testing/selftests/bpf/progs/local_storage_bench.c

   1 // SPDX-License-Identifier: GPL-2.0
   2 /* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
   3
   4 #include "vmlinux.h"
   5 #include <bpf/bpf_helpers.h>
   6 #include "bpf_misc.h"
   7
   8 #define HASHMAP_SZ 4194304
   9
  10 struct {
  11         __uint(type, BPF_MAP_TYPE_ARRAY_OF_MAPS);
  12         __uint(max_entries, 1000);
  13         __type(key, int);
  14         __type(value, int);
  15         __array(values, struct {
  16                 __uint(type, BPF_MAP_TYPE_TASK_STORAGE);
  17                 __uint(map_flags, BPF_F_NO_PREALLOC);
  18                 __type(key, int);
  19                 __type(value, int);
  20         });
  21 } array_of_local_storage_maps SEC(".maps");
  22
  23 struct {
  24         __uint(type, BPF_MAP_TYPE_ARRAY_OF_MAPS);
  25         __uint(max_entries, 1000);
  26         __type(key, int);
  27         __type(value, int);
  28         __array(values, struct {
  29                 __uint(type, BPF_MAP_TYPE_HASH);
  30                 __uint(max_entries, HASHMAP_SZ);
  31                 __type(key, int);
  32                 __type(value, int);
  33         });
  34 } array_of_hash_maps SEC(".maps");
  35
  36 long important_hits;
  37 long hits;
  38
  39 /* set from user-space */
  40 const volatile unsigned int use_hashmap;
  41 const volatile unsigned int hashmap_num_keys;
  42 const volatile unsigned int num_maps;
  43 const volatile unsigned int interleave;
  44
  45 struct loop_ctx {
  46         struct task_struct *task;
  47         long loop_hits;
  48         long loop_important_hits;
  49 };
  50
  51 static int do_lookup(unsigned int elem, struct loop_ctx *lctx)
  52 {
  53         void *map, *inner_map;
  54         int idx = 0;
  55
  56         if (use_hashmap)
  57                 map = &array_of_hash_maps;
  58         else
  59                 map = &array_of_local_storage_maps;
  60
  61         inner_map = bpf_map_lookup_elem(map, &elem);
  62         if (!inner_map)
  63                 return -1;
  64
  65         if (use_hashmap) {
  66                 idx = bpf_get_prandom_u32() % hashmap_num_keys;
  67                 bpf_map_lookup_elem(inner_map, &idx);
  68         } else {
  69                 bpf_task_storage_get(inner_map, lctx->task, &idx,
  70                                      BPF_LOCAL_STORAGE_GET_F_CREATE);
  71         }
  72
  73         lctx->loop_hits++;
  74         if (!elem)
  75                 lctx->loop_important_hits++;
  76         return 0;
  77 }
  78
  79 static long loop(u32 index, void *ctx)
  80 {
  81         struct loop_ctx *lctx = (struct loop_ctx *)ctx;
  82         unsigned int map_idx = index % num_maps;
  83
  84         do_lookup(map_idx, lctx);
  85         if (interleave && map_idx % 3 == 0)
  86                 do_lookup(0, lctx);
  87         return 0;
  88 }
  89
  90 SEC("fentry/" SYS_PREFIX "sys_getpgid")
  91 int get_local(void *ctx)
  92 {
  93         struct loop_ctx lctx;
  94
  95         lctx.task = bpf_get_current_task_btf();
  96         lctx.loop_hits = 0;
  97         lctx.loop_important_hits = 0;
  98         bpf_loop(10000, &loop, &lctx, 0);
  99         __sync_add_and_fetch(&hits, lctx.loop_hits);
 100         __sync_add_and_fetch(&important_hits, lctx.loop_important_hits);
 101         return 0;
 102 }
 103
 104 char _license[] SEC("license") = "GPL";