xref: /linux/tools/testing/selftests/bpf/progs/local_storage_bench.c (revision be239684b18e1cdcafcf8c7face4a2f562c745ad)
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";
105