1// SPDX-License-Identifier: GPL-2.0
2/* Copyright (c) 2024 Benjamin Tissoires
3 */
4
5#include "bpf_experimental.h"
6#include <bpf/bpf_helpers.h>
7#include "bpf_misc.h"
8#include "../bpf_testmod/bpf_testmod_kfunc.h"
9
10char _license[] SEC("license") = "GPL";
11
12struct hmap_elem {
13	int counter;
14	struct bpf_timer timer; /* unused */
15	struct bpf_spin_lock lock; /* unused */
16	struct bpf_wq work;
17};
18
19struct {
20	__uint(type, BPF_MAP_TYPE_HASH);
21	__uint(max_entries, 1000);
22	__type(key, int);
23	__type(value, struct hmap_elem);
24} hmap SEC(".maps");
25
26struct {
27	__uint(type, BPF_MAP_TYPE_HASH);
28	__uint(map_flags, BPF_F_NO_PREALLOC);
29	__uint(max_entries, 1000);
30	__type(key, int);
31	__type(value, struct hmap_elem);
32} hmap_malloc SEC(".maps");
33
34struct elem {
35	struct bpf_wq w;
36};
37
38struct {
39	__uint(type, BPF_MAP_TYPE_ARRAY);
40	__uint(max_entries, 2);
41	__type(key, int);
42	__type(value, struct elem);
43} array SEC(".maps");
44
45struct {
46	__uint(type, BPF_MAP_TYPE_LRU_HASH);
47	__uint(max_entries, 4);
48	__type(key, int);
49	__type(value, struct elem);
50} lru SEC(".maps");
51
52__u32 ok;
53__u32 ok_sleepable;
54
55static int test_elem_callback(void *map, int *key,
56		int (callback_fn)(void *map, int *key, struct bpf_wq *wq))
57{
58	struct elem init = {}, *val;
59	struct bpf_wq *wq;
60
61	if ((ok & (1 << *key) ||
62	    (ok_sleepable & (1 << *key))))
63		return -22;
64
65	if (map == &lru &&
66	    bpf_map_update_elem(map, key, &init, 0))
67		return -1;
68
69	val = bpf_map_lookup_elem(map, key);
70	if (!val)
71		return -2;
72
73	wq = &val->w;
74	if (bpf_wq_init(wq, map, 0) != 0)
75		return -3;
76
77	if (bpf_wq_set_callback(wq, callback_fn, 0))
78		return -4;
79
80	if (bpf_wq_start(wq, 0))
81		return -5;
82
83	return 0;
84}
85
86static int test_hmap_elem_callback(void *map, int *key,
87		int (callback_fn)(void *map, int *key, struct bpf_wq *wq))
88{
89	struct hmap_elem init = {}, *val;
90	struct bpf_wq *wq;
91
92	if ((ok & (1 << *key) ||
93	    (ok_sleepable & (1 << *key))))
94		return -22;
95
96	if (bpf_map_update_elem(map, key, &init, 0))
97		return -1;
98
99	val = bpf_map_lookup_elem(map, key);
100	if (!val)
101		return -2;
102
103	wq = &val->work;
104	if (bpf_wq_init(wq, map, 0) != 0)
105		return -3;
106
107	if (bpf_wq_set_callback(wq, callback_fn, 0))
108		return -4;
109
110	if (bpf_wq_start(wq, 0))
111		return -5;
112
113	return 0;
114}
115
116/* callback for non sleepable workqueue */
117static int wq_callback(void *map, int *key, struct bpf_wq *work)
118{
119	bpf_kfunc_common_test();
120	ok |= (1 << *key);
121	return 0;
122}
123
124/* callback for sleepable workqueue */
125static int wq_cb_sleepable(void *map, int *key, struct bpf_wq *work)
126{
127	bpf_kfunc_call_test_sleepable();
128	ok_sleepable |= (1 << *key);
129	return 0;
130}
131
132SEC("tc")
133/* test that workqueues can be used from an array */
134__retval(0)
135long test_call_array_sleepable(void *ctx)
136{
137	int key = 0;
138
139	return test_elem_callback(&array, &key, wq_cb_sleepable);
140}
141
142SEC("syscall")
143/* Same test than above but from a sleepable context. */
144__retval(0)
145long test_syscall_array_sleepable(void *ctx)
146{
147	int key = 1;
148
149	return test_elem_callback(&array, &key, wq_cb_sleepable);
150}
151
152SEC("tc")
153/* test that workqueues can be used from a hashmap */
154__retval(0)
155long test_call_hash_sleepable(void *ctx)
156{
157	int key = 2;
158
159	return test_hmap_elem_callback(&hmap, &key, wq_callback);
160}
161
162SEC("tc")
163/* test that workqueues can be used from a hashmap with NO_PREALLOC. */
164__retval(0)
165long test_call_hash_malloc_sleepable(void *ctx)
166{
167	int key = 3;
168
169	return test_hmap_elem_callback(&hmap_malloc, &key, wq_callback);
170}
171
172SEC("tc")
173/* test that workqueues can be used from a LRU map */
174__retval(0)
175long test_call_lru_sleepable(void *ctx)
176{
177	int key = 4;
178
179	return test_elem_callback(&lru, &key, wq_callback);
180}
181