bpf/progs/local_storage_bench.c

0001 // SPDX-License-Identifier: GPL-2.0
0002 /* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
0003
0004 #include "vmlinux.h"
0005 #include <bpf/bpf_helpers.h>
0006 #include "bpf_misc.h"
0007
0008 #define HASHMAP_SZ 4194304
0009
0010 struct {
0011     __uint(type, BPF_MAP_TYPE_ARRAY_OF_MAPS);
0012     __uint(max_entries, 1000);
0013     __type(key, int);
0014     __type(value, int);
0015     __array(values, struct {
0016         __uint(type, BPF_MAP_TYPE_TASK_STORAGE);
0017         __uint(map_flags, BPF_F_NO_PREALLOC);
0018         __type(key, int);
0019         __type(value, int);
0020     });
0021 } array_of_local_storage_maps SEC(".maps");
0022
0023 struct {
0024     __uint(type, BPF_MAP_TYPE_ARRAY_OF_MAPS);
0025     __uint(max_entries, 1000);
0026     __type(key, int);
0027     __type(value, int);
0028     __array(values, struct {
0029         __uint(type, BPF_MAP_TYPE_HASH);
0030         __uint(max_entries, HASHMAP_SZ);
0031         __type(key, int);
0032         __type(value, int);
0033     });
0034 } array_of_hash_maps SEC(".maps");
0035
0036 long important_hits;
0037 long hits;
0038
0039 /* set from user-space */
0040 const volatile unsigned int use_hashmap;
0041 const volatile unsigned int hashmap_num_keys;
0042 const volatile unsigned int num_maps;
0043 const volatile unsigned int interleave;
0044
0045 struct loop_ctx {
0046     struct task_struct *task;
0047     long loop_hits;
0048     long loop_important_hits;
0049 };
0050
0051 static int do_lookup(unsigned int elem, struct loop_ctx *lctx)
0052 {
0053     void *map, *inner_map;
0054     int idx = 0;
0055
0056     if (use_hashmap)
0057         map = &array_of_hash_maps;
0058     else
0059         map = &array_of_local_storage_maps;
0060
0061     inner_map = bpf_map_lookup_elem(map, &elem);
0062     if (!inner_map)
0063         return -1;
0064
0065     if (use_hashmap) {
0066         idx = bpf_get_prandom_u32() % hashmap_num_keys;
0067         bpf_map_lookup_elem(inner_map, &idx);
0068     } else {
0069         bpf_task_storage_get(inner_map, lctx->task, &idx,
0070                      BPF_LOCAL_STORAGE_GET_F_CREATE);
0071     }
0072
0073     lctx->loop_hits++;
0074     if (!elem)
0075         lctx->loop_important_hits++;
0076     return 0;
0077 }
0078
0079 static long loop(u32 index, void *ctx)
0080 {
0081     struct loop_ctx *lctx = (struct loop_ctx *)ctx;
0082     unsigned int map_idx = index % num_maps;
0083
0084     do_lookup(map_idx, lctx);
0085     if (interleave && map_idx % 3 == 0)
0086         do_lookup(0, lctx);
0087     return 0;
0088 }
0089
0090 SEC("fentry/" SYS_PREFIX "sys_getpgid")
0091 int get_local(void *ctx)
0092 {
0093     struct loop_ctx lctx;
0094
0095     lctx.task = bpf_get_current_task_btf();
0096     lctx.loop_hits = 0;
0097     lctx.loop_important_hits = 0;
0098     bpf_loop(10000, &loop, &lctx, 0);
0099     __sync_add_and_fetch(&hits, lctx.loop_hits);
0100     __sync_add_and_fetch(&important_hits, lctx.loop_important_hits);
0101     return 0;
0102 }
0103
0104 char _license[] SEC("license") = "GPL";