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	int ok_offset;
 36	struct bpf_wq w;
 37};
 38
 39struct {
 40	__uint(type, BPF_MAP_TYPE_ARRAY);
 41	__uint(max_entries, 2);
 42	__type(key, int);
 43	__type(value, struct elem);
 44} array SEC(".maps");
 45
 46struct {
 47	__uint(type, BPF_MAP_TYPE_LRU_HASH);
 48	__uint(max_entries, 4);
 49	__type(key, int);
 50	__type(value, struct elem);
 51} lru SEC(".maps");
 52
 53__u32 ok;
 54__u32 ok_sleepable;
 55
 56static int test_elem_callback(void *map, int *key,
 57		int (callback_fn)(void *map, int *key, void *value))
 58{
 59	struct elem init = {}, *val;
 60	struct bpf_wq *wq;
 61
 62	if ((ok & (1 << *key) ||
 63	    (ok_sleepable & (1 << *key))))
 64		return -22;
 65
 66	if (map == &lru &&
 67	    bpf_map_update_elem(map, key, &init, 0))
 68		return -1;
 69
 70	val = bpf_map_lookup_elem(map, key);
 71	if (!val)
 72		return -2;
 73
 74	val->ok_offset = *key;
 75
 76	wq = &val->w;
 77	if (bpf_wq_init(wq, map, 0) != 0)
 78		return -3;
 79
 80	if (bpf_wq_set_callback(wq, callback_fn, 0))
 81		return -4;
 82
 83	if (bpf_wq_start(wq, 0))
 84		return -5;
 85
 86	return 0;
 87}
 88
 89static int test_hmap_elem_callback(void *map, int *key,
 90		int (callback_fn)(void *map, int *key, void *value))
 91{
 92	struct hmap_elem init = {}, *val;
 93	struct bpf_wq *wq;
 94
 95	if ((ok & (1 << *key) ||
 96	    (ok_sleepable & (1 << *key))))
 97		return -22;
 98
 99	if (bpf_map_update_elem(map, key, &init, 0))
100		return -1;
101
102	val = bpf_map_lookup_elem(map, key);
103	if (!val)
104		return -2;
105
106	wq = &val->work;
107	if (bpf_wq_init(wq, map, 0) != 0)
108		return -3;
109
110	if (bpf_wq_set_callback(wq, callback_fn, 0))
111		return -4;
112
113	if (bpf_wq_start(wq, 0))
114		return -5;
115
116	return 0;
117}
118
119/* callback for non sleepable workqueue */
120static int wq_callback(void *map, int *key, void *value)
121{
122	bpf_kfunc_common_test();
123	ok |= (1 << *key);
124	return 0;
125}
126
127/* callback for sleepable workqueue */
128static int wq_cb_sleepable(void *map, int *key, void *value)
129{
130	struct elem *data = (struct elem *)value;
131	int offset = data->ok_offset;
132
133	if (*key != offset)
134		return 0;
135
136	bpf_kfunc_call_test_sleepable();
137	ok_sleepable |= (1 << offset);
138	return 0;
139}
140
141SEC("tc")
142/* test that workqueues can be used from an array */
143__retval(0)
144long test_call_array_sleepable(void *ctx)
145{
146	int key = 0;
147
148	return test_elem_callback(&array, &key, wq_cb_sleepable);
149}
150
151SEC("syscall")
152/* Same test than above but from a sleepable context. */
153__retval(0)
154long test_syscall_array_sleepable(void *ctx)
155{
156	int key = 1;
157
158	return test_elem_callback(&array, &key, wq_cb_sleepable);
159}
160
161SEC("tc")
162/* test that workqueues can be used from a hashmap */
163__retval(0)
164long test_call_hash_sleepable(void *ctx)
165{
166	int key = 2;
167
168	return test_hmap_elem_callback(&hmap, &key, wq_callback);
169}
170
171SEC("tc")
172/* test that workqueues can be used from a hashmap with NO_PREALLOC. */
173__retval(0)
174long test_call_hash_malloc_sleepable(void *ctx)
175{
176	int key = 3;
177
178	return test_hmap_elem_callback(&hmap_malloc, &key, wq_callback);
179}
180
181SEC("tc")
182/* test that workqueues can be used from a LRU map */
183__retval(0)
184long test_call_lru_sleepable(void *ctx)
185{
186	int key = 4;
187
188	return test_elem_callback(&lru, &key, wq_callback);
189}