blob: 71ad825a29ef03d025ae4d050748420e2c9fef53
1 | /* vi: set sw=4 ts=4: */ |
2 | /* |
3 | * taskset - retrieve or set a processes' CPU affinity |
4 | * Copyright (c) 2006 Bernhard Reutner-Fischer |
5 | * |
6 | * Licensed under GPLv2 or later, see file LICENSE in this source tree. |
7 | */ |
8 | |
9 | //usage:#define taskset_trivial_usage |
10 | //usage: "[-p] [MASK] [PID | PROG ARGS]" |
11 | //usage:#define taskset_full_usage "\n\n" |
12 | //usage: "Set or get CPU affinity\n" |
13 | //usage: "\n -p Operate on an existing PID" |
14 | //usage: |
15 | //usage:#define taskset_example_usage |
16 | //usage: "$ taskset 0x7 ./dgemm_test&\n" |
17 | //usage: "$ taskset -p 0x1 $!\n" |
18 | //usage: "pid 4790's current affinity mask: 7\n" |
19 | //usage: "pid 4790's new affinity mask: 1\n" |
20 | //usage: "$ taskset 0x7 /bin/sh -c './taskset -p 0x1 $$'\n" |
21 | //usage: "pid 6671's current affinity mask: 1\n" |
22 | //usage: "pid 6671's new affinity mask: 1\n" |
23 | //usage: "$ taskset -p 1\n" |
24 | //usage: "pid 1's current affinity mask: 3\n" |
25 | |
26 | #ifdef BIONIC_ICS |
27 | #define _GNU_SOURCE 1 |
28 | #endif |
29 | |
30 | #include <sched.h> |
31 | #include "libbb.h" |
32 | |
33 | #if ENABLE_FEATURE_TASKSET_FANCY |
34 | #define TASKSET_PRINTF_MASK "%s" |
35 | /* craft a string from the mask */ |
36 | static char *from_cpuset(cpu_set_t *mask) |
37 | { |
38 | int i; |
39 | char *ret = NULL; |
40 | char *str = xzalloc((CPU_SETSIZE / 4) + 1); /* we will leak it */ |
41 | |
42 | for (i = CPU_SETSIZE - 4; i >= 0; i -= 4) { |
43 | int val = 0; |
44 | int off; |
45 | for (off = 0; off <= 3; ++off) |
46 | if (CPU_ISSET(i + off, mask)) |
47 | val |= 1 << off; |
48 | if (!ret && val) |
49 | ret = str; |
50 | *str++ = bb_hexdigits_upcase[val] | 0x20; |
51 | } |
52 | return ret; |
53 | } |
54 | #else |
55 | #define TASKSET_PRINTF_MASK "%llx" |
56 | static unsigned long long from_cpuset(cpu_set_t *mask) |
57 | { |
58 | struct BUG_CPU_SETSIZE_is_too_small { |
59 | char BUG_CPU_SETSIZE_is_too_small[ |
60 | CPU_SETSIZE < sizeof(int) ? -1 : 1]; |
61 | }; |
62 | char *p = (void*)mask; |
63 | |
64 | /* Take the least significant bits. Careful! |
65 | * Consider both CPU_SETSIZE=4 and CPU_SETSIZE=1024 cases |
66 | */ |
67 | #if BB_BIG_ENDIAN |
68 | /* For big endian, it means LAST bits */ |
69 | if (CPU_SETSIZE < sizeof(long)) |
70 | p += CPU_SETSIZE - sizeof(int); |
71 | else if (CPU_SETSIZE < sizeof(long long)) |
72 | p += CPU_SETSIZE - sizeof(long); |
73 | else |
74 | p += CPU_SETSIZE - sizeof(long long); |
75 | #endif |
76 | if (CPU_SETSIZE < sizeof(long)) |
77 | return *(unsigned*)p; |
78 | if (CPU_SETSIZE < sizeof(long long)) |
79 | return *(unsigned long*)p; |
80 | return *(unsigned long long*)p; |
81 | } |
82 | #endif |
83 | |
84 | |
85 | int taskset_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE; |
86 | int taskset_main(int argc UNUSED_PARAM, char **argv) |
87 | { |
88 | cpu_set_t mask; |
89 | pid_t pid = 0; |
90 | unsigned opt_p; |
91 | const char *current_new; |
92 | char *pid_str; |
93 | char *aff = NULL; |
94 | |
95 | /* NB: we mimic util-linux's taskset: -p does not take |
96 | * an argument, i.e., "-pN" is NOT valid, only "-p N"! |
97 | * Indeed, util-linux-2.13-pre7 uses: |
98 | * getopt_long(argc, argv, "+pchV", ...), not "...p:..." */ |
99 | |
100 | opt_complementary = "-1"; /* at least 1 arg */ |
101 | opt_p = getopt32(argv, "+p"); |
102 | argv += optind; |
103 | |
104 | if (opt_p) { |
105 | pid_str = *argv++; |
106 | if (*argv) { /* "-p <aff> <pid> ...rest.is.ignored..." */ |
107 | aff = pid_str; |
108 | pid_str = *argv; /* NB: *argv != NULL in this case */ |
109 | } |
110 | /* else it was just "-p <pid>", and *argv == NULL */ |
111 | pid = xatoul_range(pid_str, 1, ((unsigned)(pid_t)ULONG_MAX) >> 1); |
112 | } else { |
113 | aff = *argv++; /* <aff> <cmd...> */ |
114 | if (!*argv) |
115 | bb_show_usage(); |
116 | } |
117 | |
118 | current_new = "current\0new"; |
119 | if (opt_p) { |
120 | print_aff: |
121 | if (sched_getaffinity(pid, sizeof(mask), &mask) < 0) |
122 | bb_perror_msg_and_die("can't %cet pid %d's affinity", 'g', pid); |
123 | printf("pid %d's %s affinity mask: "TASKSET_PRINTF_MASK"\n", |
124 | pid, current_new, from_cpuset(&mask)); |
125 | if (!*argv) { |
126 | /* Either it was just "-p <pid>", |
127 | * or it was "-p <aff> <pid>" and we came here |
128 | * for the second time (see goto below) */ |
129 | return EXIT_SUCCESS; |
130 | } |
131 | *argv = NULL; |
132 | current_new += 8; /* "new" */ |
133 | } |
134 | |
135 | { /* Affinity was specified, translate it into cpu_set_t */ |
136 | unsigned i; |
137 | /* Do not allow zero mask: */ |
138 | unsigned long long m = xstrtoull_range(aff, 0, 1, ULLONG_MAX); |
139 | enum { CNT_BIT = CPU_SETSIZE < sizeof(m)*8 ? CPU_SETSIZE : sizeof(m)*8 }; |
140 | |
141 | CPU_ZERO(&mask); |
142 | for (i = 0; i < CNT_BIT; i++) { |
143 | unsigned long long bit = (1ULL << i); |
144 | if (bit & m) |
145 | CPU_SET(i, &mask); |
146 | } |
147 | } |
148 | |
149 | /* Set pid's or our own (pid==0) affinity */ |
150 | if (sched_setaffinity(pid, sizeof(mask), &mask)) |
151 | bb_perror_msg_and_die("can't %cet pid %d's affinity", 's', pid); |
152 | |
153 | if (!argv[0]) /* "-p <aff> <pid> [...ignored...]" */ |
154 | goto print_aff; /* print new affinity and exit */ |
155 | |
156 | BB_EXECVP_or_die(argv); |
157 | } |
158 |