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 <signal.h> 8 #include <string.h> 9 #include <stdio.h> 10 #include <unistd.h> 11 #include <sched.h> 12 #include <bpf/bpf.h> 13 #include <scx/common.h> 14 #include <sys/wait.h> 15 #include "scx_test.h" 16 #include "init_enable_count.bpf.skel.h" 17 18 #define SCHED_EXT 7 19 20 static enum scx_test_status run_test(bool global) 21 { 22 struct init_enable_count *skel; 23 struct bpf_link *link; 24 const u32 num_children = 5, num_pre_forks = 1024; 25 int ret, i, status; 26 struct sched_param param = {}; 27 pid_t pids[num_pre_forks]; 28 int pipe_fds[2]; 29 30 SCX_FAIL_IF(pipe(pipe_fds) < 0, "Failed to create pipe"); 31 32 skel = init_enable_count__open(); 33 SCX_FAIL_IF(!skel, "Failed to open"); 34 SCX_ENUM_INIT(skel); 35 36 if (!global) 37 skel->struct_ops.init_enable_count_ops->flags |= SCX_OPS_SWITCH_PARTIAL; 38 39 SCX_FAIL_IF(init_enable_count__load(skel), "Failed to load skel"); 40 41 /* 42 * Fork a bunch of children before we attach the scheduler so that we 43 * ensure (at least in practical terms) that there are more tasks that 44 * transition from SCHED_OTHER -> SCHED_EXT than there are tasks that 45 * take the fork() path either below or in other processes. 46 * 47 * All children will block on read() on the pipe until the parent closes 48 * the write end after attaching the scheduler, which signals all of 49 * them to exit simultaneously. Auto-reap so we don't have to wait on 50 * them. 51 */ 52 signal(SIGCHLD, SIG_IGN); 53 for (i = 0; i < num_pre_forks; i++) { 54 pid_t pid = fork(); 55 56 SCX_FAIL_IF(pid < 0, "Failed to fork child"); 57 if (pid == 0) { 58 char buf; 59 60 close(pipe_fds[1]); 61 if (read(pipe_fds[0], &buf, 1) < 0) 62 exit(1); 63 close(pipe_fds[0]); 64 exit(0); 65 } 66 } 67 close(pipe_fds[0]); 68 69 link = bpf_map__attach_struct_ops(skel->maps.init_enable_count_ops); 70 SCX_FAIL_IF(!link, "Failed to attach struct_ops"); 71 72 /* Signal all pre-forked children to exit. */ 73 close(pipe_fds[1]); 74 signal(SIGCHLD, SIG_DFL); 75 76 SCX_ASSERT_ALIVE(skel, uei); 77 78 bpf_link__destroy(link); 79 80 /* 81 * The clean unregister above is recorded in uei and would 82 * otherwise trip the liveness check of the second cycle, which 83 * reuses this skeleton. 84 */ 85 memset(&skel->data->uei, 0, sizeof(skel->data->uei)); 86 87 SCX_GE(skel->bss->init_task_cnt, num_pre_forks); 88 SCX_GE(skel->bss->exit_task_cnt, num_pre_forks); 89 90 link = bpf_map__attach_struct_ops(skel->maps.init_enable_count_ops); 91 SCX_FAIL_IF(!link, "Failed to attach struct_ops"); 92 93 /* SCHED_EXT children */ 94 for (i = 0; i < num_children; i++) { 95 pids[i] = fork(); 96 SCX_FAIL_IF(pids[i] < 0, "Failed to fork child"); 97 98 if (pids[i] == 0) { 99 ret = sched_setscheduler(0, SCHED_EXT, ¶m); 100 SCX_BUG_ON(ret, "Failed to set sched to sched_ext"); 101 102 /* 103 * Reset to SCHED_OTHER for half of them. Counts for 104 * everything should still be the same regardless, as 105 * ops.disable() is invoked even if a task is still on 106 * SCHED_EXT before it exits. 107 */ 108 if (i % 2 == 0) { 109 ret = sched_setscheduler(0, SCHED_OTHER, ¶m); 110 SCX_BUG_ON(ret, "Failed to reset sched to normal"); 111 } 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 SCX child\n"); 118 119 SCX_FAIL_IF(status != 0, "SCX child %d exited with status %d\n", i, 120 status); 121 } 122 123 /* SCHED_OTHER children */ 124 for (i = 0; i < num_children; i++) { 125 pids[i] = fork(); 126 if (pids[i] == 0) 127 exit(0); 128 } 129 130 for (i = 0; i < num_children; i++) { 131 SCX_FAIL_IF(waitpid(pids[i], &status, 0) != pids[i], 132 "Failed to wait for normal child\n"); 133 134 SCX_FAIL_IF(status != 0, "Normal child %d exited with status %d\n", i, 135 status); 136 } 137 138 SCX_ASSERT_ALIVE(skel, uei); 139 140 bpf_link__destroy(link); 141 142 SCX_GE(skel->bss->init_task_cnt, 2 * num_children); 143 SCX_GE(skel->bss->exit_task_cnt, 2 * num_children); 144 145 if (global) { 146 SCX_GE(skel->bss->enable_cnt, 2 * num_children); 147 SCX_GE(skel->bss->disable_cnt, 2 * num_children); 148 } else { 149 SCX_EQ(skel->bss->enable_cnt, num_children); 150 SCX_EQ(skel->bss->disable_cnt, num_children); 151 } 152 /* 153 * We forked a ton of tasks before we attached the scheduler above, so 154 * this should be fine. Technically it could be flaky if a ton of forks 155 * are happening at the same time in other processes, but that should 156 * be exceedingly unlikely. 157 */ 158 SCX_GT(skel->bss->init_transition_cnt, skel->bss->init_fork_cnt); 159 SCX_GE(skel->bss->init_fork_cnt, 2 * num_children); 160 161 init_enable_count__destroy(skel); 162 163 return SCX_TEST_PASS; 164 } 165 166 static enum scx_test_status run(void *ctx) 167 { 168 enum scx_test_status status; 169 170 status = run_test(true); 171 if (status != SCX_TEST_PASS) 172 return status; 173 174 return run_test(false); 175 } 176 177 struct scx_test init_enable_count = { 178 .name = "init_enable_count", 179 .description = "Verify we correctly count the occurrences of init, " 180 "enable, etc callbacks.", 181 .run = run, 182 }; 183 REGISTER_SCX_TEST(&init_enable_count) 184