1 /* SPDX-License-Identifier: GPL-2.0 */
2 /*
3 * Copyright (c) 2023 Meta Platforms, Inc. and affiliates.
4 * Copyright (c) 2023 David Vernet <dvernet@meta.com>
5 * Copyright (c) 2023 Tejun Heo <tj@kernel.org>
6 */
7 #include <stdio.h>
8 #include <unistd.h>
9 #include <sched.h>
10 #include <bpf/bpf.h>
11 #include <scx/common.h>
12 #include <sys/wait.h>
13 #include "scx_test.h"
14 #include "init_enable_count.bpf.skel.h"
15
16 #define SCHED_EXT 7
17
18 static struct init_enable_count *
19 open_load_prog(bool global)
20 {
21 struct init_enable_count *skel;
22
23 skel = init_enable_count__open();
24 SCX_BUG_ON(!skel, "Failed to open skel");
25
26 if (!global)
27 skel->struct_ops.init_enable_count_ops->flags |= SCX_OPS_SWITCH_PARTIAL;
28
29 SCX_BUG_ON(init_enable_count__load(skel), "Failed to load skel");
30
31 return skel;
32 }
33
34 static enum scx_test_status run_test(bool global)
35 {
36 struct init_enable_count *skel;
37 struct bpf_link *link;
38 const u32 num_children = 5, num_pre_forks = 1024;
39 int ret, i, status;
40 struct sched_param param = {};
41 pid_t pids[num_pre_forks];
42
43 skel = open_load_prog(global);
44
45 /*
46 * Fork a bunch of children before we attach the scheduler so that we
47 * ensure (at least in practical terms) that there are more tasks that
48 * transition from SCHED_OTHER -> SCHED_EXT than there are tasks that
49 * take the fork() path either below or in other processes.
50 */
51 for (i = 0; i < num_pre_forks; i++) {
52 pids[i] = fork();
53 SCX_FAIL_IF(pids[i] < 0, "Failed to fork child");
54 if (pids[i] == 0) {
55 sleep(1);
56 exit(0);
57 }
58 }
59
60 link = bpf_map__attach_struct_ops(skel->maps.init_enable_count_ops);
61 SCX_FAIL_IF(!link, "Failed to attach struct_ops");
62
63 for (i = 0; i < num_pre_forks; i++) {
64 SCX_FAIL_IF(waitpid(pids[i], &status, 0) != pids[i],
65 "Failed to wait for pre-forked child\n");
66
67 SCX_FAIL_IF(status != 0, "Pre-forked child %d exited with status %d\n", i,
68 status);
69 }
70
71 bpf_link__destroy(link);
72 SCX_GE(skel->bss->init_task_cnt, num_pre_forks);
73 SCX_GE(skel->bss->exit_task_cnt, num_pre_forks);
74
75 link = bpf_map__attach_struct_ops(skel->maps.init_enable_count_ops);
76 SCX_FAIL_IF(!link, "Failed to attach struct_ops");
77
78 /* SCHED_EXT children */
79 for (i = 0; i < num_children; i++) {
80 pids[i] = fork();
81 SCX_FAIL_IF(pids[i] < 0, "Failed to fork child");
82
83 if (pids[i] == 0) {
84 ret = sched_setscheduler(0, SCHED_EXT, ¶m);
85 SCX_BUG_ON(ret, "Failed to set sched to sched_ext");
86
87 /*
88 * Reset to SCHED_OTHER for half of them. Counts for
89 * everything should still be the same regardless, as
90 * ops.disable() is invoked even if a task is still on
91 * SCHED_EXT before it exits.
92 */
93 if (i % 2 == 0) {
94 ret = sched_setscheduler(0, SCHED_OTHER, ¶m);
95 SCX_BUG_ON(ret, "Failed to reset sched to normal");
96 }
97 exit(0);
98 }
99 }
100 for (i = 0; i < num_children; i++) {
101 SCX_FAIL_IF(waitpid(pids[i], &status, 0) != pids[i],
102 "Failed to wait for SCX child\n");
103
104 SCX_FAIL_IF(status != 0, "SCX child %d exited with status %d\n", i,
105 status);
106 }
107
108 /* SCHED_OTHER children */
109 for (i = 0; i < num_children; i++) {
110 pids[i] = fork();
111 if (pids[i] == 0)
112 exit(0);
113 }
114
115 for (i = 0; i < num_children; i++) {
116 SCX_FAIL_IF(waitpid(pids[i], &status, 0) != pids[i],
117 "Failed to wait for normal child\n");
118
119 SCX_FAIL_IF(status != 0, "Normal child %d exited with status %d\n", i,
120 status);
121 }
122
123 bpf_link__destroy(link);
124
125 SCX_GE(skel->bss->init_task_cnt, 2 * num_children);
126 SCX_GE(skel->bss->exit_task_cnt, 2 * num_children);
127
128 if (global) {
129 SCX_GE(skel->bss->enable_cnt, 2 * num_children);
130 SCX_GE(skel->bss->disable_cnt, 2 * num_children);
131 } else {
132 SCX_EQ(skel->bss->enable_cnt, num_children);
133 SCX_EQ(skel->bss->disable_cnt, num_children);
134 }
135 /*
136 * We forked a ton of tasks before we attached the scheduler above, so
137 * this should be fine. Technically it could be flaky if a ton of forks
138 * are happening at the same time in other processes, but that should
139 * be exceedingly unlikely.
140 */
141 SCX_GT(skel->bss->init_transition_cnt, skel->bss->init_fork_cnt);
142 SCX_GE(skel->bss->init_fork_cnt, 2 * num_children);
143
144 init_enable_count__destroy(skel);
145
146 return SCX_TEST_PASS;
147 }
148
149 static enum scx_test_status run(void *ctx)
150 {
151 enum scx_test_status status;
152
153 status = run_test(true);
154 if (status != SCX_TEST_PASS)
155 return status;
156
157 return run_test(false);
158 }
159
160 struct scx_test init_enable_count = {
161 .name = "init_enable_count",
162 .description = "Verify we do the correct amount of counting of init, "
163 "enable, etc callbacks.",
164 .run = run,
165 };
166 REGISTER_SCX_TEST(&init_enable_count)
167
Linux® is a registered trademark of Linus Torvalds in the United States and other countries.
TOMOYO® is a registered trademark of NTT DATA CORPORATION.