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