1// SPDX-License-Identifier: GPL-2.0
2/* Copyright (c) 2021 Facebook */
3
4#include <errno.h>
5#include <linux/bpf.h>
6#include <stdbool.h>
7#include <bpf/bpf_helpers.h>
8#include "bpf_misc.h"
9
10char _license[] SEC("license") = "GPL";
11
12struct bpf_map;
13
14__u8 rand_vals[2500000];
15const __u32 nr_rand_bytes = 2500000;
16
17struct {
18	__uint(type, BPF_MAP_TYPE_ARRAY);
19	__uint(key_size, sizeof(__u32));
20	/* max entries and value_size will be set programmatically.
21	 * They are configurable from the userspace bench program.
22	 */
23} array_map SEC(".maps");
24
25struct {
26	__uint(type, BPF_MAP_TYPE_BLOOM_FILTER);
27	/* max entries,  value_size, and # of hash functions will be set
28	 * programmatically. They are configurable from the userspace
29	 * bench program.
30	 */
31	__uint(map_extra, 3);
32} bloom_map SEC(".maps");
33
34struct {
35	__uint(type, BPF_MAP_TYPE_HASH);
36	/* max entries, key_size, and value_size, will be set
37	 * programmatically. They are configurable from the userspace
38	 * bench program.
39	 */
40} hashmap SEC(".maps");
41
42struct callback_ctx {
43	struct bpf_map *map;
44	bool update;
45};
46
47/* Tracks the number of hits, drops, and false hits */
48struct {
49	__u32 stats[3];
50} __attribute__((__aligned__(256))) percpu_stats[256];
51
52const __u32 hit_key  = 0;
53const __u32 drop_key  = 1;
54const __u32 false_hit_key = 2;
55
56__u8 value_size;
57
58const volatile bool hashmap_use_bloom;
59const volatile bool count_false_hits;
60
61int error = 0;
62
63static __always_inline void log_result(__u32 key)
64{
65	__u32 cpu = bpf_get_smp_processor_id();
66
67	percpu_stats[cpu & 255].stats[key]++;
68}
69
70static __u64
71bloom_callback(struct bpf_map *map, __u32 *key, void *val,
72	       struct callback_ctx *data)
73{
74	int err;
75
76	if (data->update)
77		err = bpf_map_push_elem(data->map, val, 0);
78	else
79		err = bpf_map_peek_elem(data->map, val);
80
81	if (err) {
82		error |= 1;
83		return 1; /* stop the iteration */
84	}
85
86	log_result(hit_key);
87
88	return 0;
89}
90
91SEC("fentry/" SYS_PREFIX "sys_getpgid")
92int bloom_lookup(void *ctx)
93{
94	struct callback_ctx data;
95
96	data.map = (struct bpf_map *)&bloom_map;
97	data.update = false;
98
99	bpf_for_each_map_elem(&array_map, bloom_callback, &data, 0);
100
101	return 0;
102}
103
104SEC("fentry/" SYS_PREFIX "sys_getpgid")
105int bloom_update(void *ctx)
106{
107	struct callback_ctx data;
108
109	data.map = (struct bpf_map *)&bloom_map;
110	data.update = true;
111
112	bpf_for_each_map_elem(&array_map, bloom_callback, &data, 0);
113
114	return 0;
115}
116
117SEC("fentry/" SYS_PREFIX "sys_getpgid")
118int bloom_hashmap_lookup(void *ctx)
119{
120	__u64 *result;
121	int i, err;
122
123	__u32 index = bpf_get_prandom_u32();
124	__u32 bitmask = (1ULL << 21) - 1;
125
126	for (i = 0; i < 1024; i++, index += value_size) {
127		index = index & bitmask;
128
129		if (hashmap_use_bloom) {
130			err = bpf_map_peek_elem(&bloom_map,
131						rand_vals + index);
132			if (err) {
133				if (err != -ENOENT) {
134					error |= 2;
135					return 0;
136				}
137				log_result(hit_key);
138				continue;
139			}
140		}
141
142		result = bpf_map_lookup_elem(&hashmap,
143					     rand_vals + index);
144		if (result) {
145			log_result(hit_key);
146		} else {
147			if (hashmap_use_bloom && count_false_hits)
148				log_result(false_hit_key);
149			log_result(drop_key);
150		}
151	}
152
153	return 0;
154}
155