xref: /linux/tools/testing/selftests/bpf/progs/task_work.c (revision 07fdad3a93756b872da7b53647715c48d0f4a2d0)
1 // SPDX-License-Identifier: GPL-2.0
2 /* Copyright (c) 2025 Meta Platforms, Inc. and affiliates. */
3 
4 #include <vmlinux.h>
5 #include <string.h>
6 #include <stdbool.h>
7 #include <bpf/bpf_helpers.h>
8 #include <bpf/bpf_tracing.h>
9 #include "bpf_misc.h"
10 #include "errno.h"
11 
12 char _license[] SEC("license") = "GPL";
13 
14 const void *user_ptr = NULL;
15 
16 struct elem {
17 	char data[128];
18 	struct bpf_task_work tw;
19 };
20 
21 struct {
22 	__uint(type, BPF_MAP_TYPE_HASH);
23 	__uint(map_flags, BPF_F_NO_PREALLOC);
24 	__uint(max_entries, 1);
25 	__type(key, int);
26 	__type(value, struct elem);
27 } hmap SEC(".maps");
28 
29 struct {
30 	__uint(type, BPF_MAP_TYPE_ARRAY);
31 	__uint(max_entries, 1);
32 	__type(key, int);
33 	__type(value, struct elem);
34 } arrmap SEC(".maps");
35 
36 struct {
37 	__uint(type, BPF_MAP_TYPE_LRU_HASH);
38 	__uint(max_entries, 1);
39 	__type(key, int);
40 	__type(value, struct elem);
41 } lrumap SEC(".maps");
42 
43 static int process_work(struct bpf_map *map, void *key, void *value)
44 {
45 	struct elem *work = value;
46 
47 	bpf_copy_from_user_str(work->data, sizeof(work->data), (const void *)user_ptr, 0);
48 	return 0;
49 }
50 
51 int key = 0;
52 
53 SEC("perf_event")
54 int oncpu_hash_map(struct pt_regs *args)
55 {
56 	struct elem empty_work = { .data = { 0 } };
57 	struct elem *work;
58 	struct task_struct *task;
59 	int err;
60 
61 	task = bpf_get_current_task_btf();
62 	err = bpf_map_update_elem(&hmap, &key, &empty_work, BPF_NOEXIST);
63 	if (err)
64 		return 0;
65 	work = bpf_map_lookup_elem(&hmap, &key);
66 	if (!work)
67 		return 0;
68 
69 	bpf_task_work_schedule_resume(task, &work->tw, &hmap, process_work, NULL);
70 	return 0;
71 }
72 
73 SEC("perf_event")
74 int oncpu_array_map(struct pt_regs *args)
75 {
76 	struct elem *work;
77 	struct task_struct *task;
78 
79 	task = bpf_get_current_task_btf();
80 	work = bpf_map_lookup_elem(&arrmap, &key);
81 	if (!work)
82 		return 0;
83 	bpf_task_work_schedule_signal(task, &work->tw, &arrmap, process_work, NULL);
84 	return 0;
85 }
86 
87 SEC("perf_event")
88 int oncpu_lru_map(struct pt_regs *args)
89 {
90 	struct elem empty_work = { .data = { 0 } };
91 	struct elem *work;
92 	struct task_struct *task;
93 	int err;
94 
95 	task = bpf_get_current_task_btf();
96 	work = bpf_map_lookup_elem(&lrumap, &key);
97 	if (work)
98 		return 0;
99 	err = bpf_map_update_elem(&lrumap, &key, &empty_work, BPF_NOEXIST);
100 	if (err)
101 		return 0;
102 	work = bpf_map_lookup_elem(&lrumap, &key);
103 	if (!work || work->data[0])
104 		return 0;
105 	bpf_task_work_schedule_resume(task, &work->tw, &lrumap, process_work, NULL);
106 	return 0;
107 }
108