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
10struct {
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
23struct {
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
36long important_hits;
37long hits;
38
39/* set from user-space */
40const volatile unsigned int use_hashmap;
41const volatile unsigned int hashmap_num_keys;
42const volatile unsigned int num_maps;
43const volatile unsigned int interleave;
44
45struct loop_ctx {
46 struct task_struct *task;
47 long loop_hits;
48 long loop_important_hits;
49};
50
51static 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
79static 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(elem: map_idx, lctx);
85 if (interleave && map_idx % 3 == 0)
86 do_lookup(elem: 0, lctx);
87 return 0;
88}
89
90SEC("fentry/" SYS_PREFIX "sys_getpgid")
91int 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
104char _license[] SEC("license") = "GPL";
105

source code of linux/tools/testing/selftests/bpf/progs/local_storage_bench.c