1 /*
2 * Copyright 2023 The ChromiumOS Authors
3 * Use of this source code is governed by a BSD-style license that can be
4 * found in the LICENSE file.
5 */
6 #include <err.h>
7 #include <errno.h>
8 #include <linux/ioprio.h>
9 #include <sched.h>
10 #include <stdio.h>
11 #include <stdlib.h>
12 #include <sys/resource.h>
13 #include <sys/syscall.h>
14 #include <sys/sysinfo.h>
15 #include <unistd.h>
16
17 #include <string_view>
18
19 static const char description[] = R"(bgtask - `nice -n19` but more.
20
21 This program exists to set a few process attributes, and exec another program.
22
23 Intended usage is:
24 $ bgtask ./my_long_running_build --extra-optimizations --and-more
25
26 Specifically, `bgtask`:
27 - sets its own priority so essentially any other task will take priority, then
28 - sets its I/O priority so any regular task will take priority, then
29 - sets its CPU mask so it may only run on 9/10ths of your cores (this step is
30 skipped if you've already got a more restrictive mask), then
31 - execs the command you gave it.
32 )";
33
print_help_and_exit(int exit_code)34 static void print_help_and_exit(int exit_code) {
35 fputs(description, stderr);
36 exit(exit_code);
37 }
38
deprioritize_nice_or_warn()39 static void deprioritize_nice_or_warn() {
40 constexpr int current_process = 0;
41 constexpr int max_nice_priority = 19;
42 if (setpriority(PRIO_PROCESS, current_process, max_nice_priority)) {
43 warn("setting priority failed");
44 }
45 }
46
deprioritize_io_or_warn()47 static void deprioritize_io_or_warn() {
48 // Glibc provides no wrapper here.
49 constexpr int current_process = 0;
50 const int background_io_prio = IOPRIO_PRIO_VALUE(IOPRIO_CLASS_IDLE, 0);
51 if (syscall(SYS_ioprio_set, IOPRIO_WHO_PROCESS, current_process,
52 background_io_prio) == -1) {
53 warn("setting ioprio failed");
54 }
55 }
56
restrict_cpu_mask_or_warn()57 static void restrict_cpu_mask_or_warn() {
58 const int available_cpus = get_nprocs();
59 // Use 9/10ths of CPUs, as requested in the mask.
60 const int cpus_to_use = (available_cpus * 9) / 10;
61 // Erm, single-core systems need not apply.
62 if (cpus_to_use == 0) {
63 return;
64 }
65
66 // Note pid==0 means "the current thread," for sched_*affinity calls.
67 constexpr int current_process = 0;
68 cpu_set_t current_mask;
69 CPU_ZERO(¤t_mask);
70 if (sched_getaffinity(current_process, sizeof(current_mask), ¤t_mask) ==
71 -1) {
72 if (errno == EINVAL) {
73 // This can only happen if a machine has >1K cores. That can be handled,
74 // but is extra complexity that I can't test & isn't expected to
75 // realistically be a problem in the next few years.
76 warnx("statically-allocated cpu affinity mask is too small");
77 } else {
78 warn("sched_getaffinity failed");
79 }
80 return;
81 }
82
83 const int cpus_in_current_mask = CPU_COUNT(¤t_mask);
84 int cpus_to_disable = cpus_in_current_mask - cpus_to_use;
85 if (cpus_to_disable <= 0) {
86 // Don't warn; this probably isn't useful informtion to the user.
87 return;
88 }
89
90 for (int i = 0; i < available_cpus; ++i) {
91 if (CPU_ISSET(i, ¤t_mask)) {
92 CPU_CLR(i, ¤t_mask);
93 --cpus_to_disable;
94 if (cpus_to_disable == 0) {
95 break;
96 }
97 }
98 }
99
100 if (cpus_to_disable != 0) {
101 warnx("Internal error: iterated through CPU mask but had %d CPUs left to "
102 "disable.",
103 cpus_to_disable);
104 return;
105 }
106
107 if (sched_setaffinity(current_process, sizeof(current_mask), ¤t_mask) ==
108 -1) {
109 warn("sched_setaffinity failed");
110 }
111 }
112
main(int argc,char ** argv)113 __attribute__((noreturn)) int main(int argc, char **argv) {
114 if (argc == 1) {
115 print_help_and_exit(/*exit_code=*/1);
116 }
117
118 std::string_view argv1 = argv[1];
119 // Do minimal option parsing, since the user is likely to be passing `-flags`
120 // and `--flags` to the program they're invoking.
121 if (argv1 == "-h" || argv1 == "--help") {
122 print_help_and_exit(/*exit_code=*/0);
123 }
124
125 deprioritize_nice_or_warn();
126 deprioritize_io_or_warn();
127 restrict_cpu_mask_or_warn();
128
129 execvp(argv[1], &argv[1]);
130 err(1, "execvp failed");
131 }
132