167 lines · c
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 717 18static struct init_enable_count *19open_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 34static 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 we47 * ensure (at least in practical terms) that there are more tasks that48 * transition from SCHED_OTHER -> SCHED_EXT than there are tasks that49 * 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 for89 * everything should still be the same regardless, as90 * ops.disable() is invoked even if a task is still on91 * 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, so137 * this should be fine. Technically it could be flaky if a ton of forks138 * are happening at the same time in other processes, but that should139 * 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 149static 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 160struct 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};166REGISTER_SCX_TEST(&init_enable_count)167