summaryrefslogtreecommitdiff
path: root/miscutils/taskset.c
blob: 77fc8643d5eba73bca5a613de9bb09db1eae80ba (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
/* vi: set sw=4 ts=4: */
/*
 * taskset - retrieve or set a processes' CPU affinity
 * Copyright (c) 2006 Bernhard Reutner-Fischer
 *
 * Licensed under GPLv2 or later, see file LICENSE in this source tree.
 */

//usage:#define taskset_trivial_usage
//usage:       "[-p] [MASK] [PID | PROG ARGS]"
//usage:#define taskset_full_usage "\n\n"
//usage:       "Set or get CPU affinity\n"
//usage:     "\nOptions:"
//usage:     "\n	-p	Operate on an existing PID"
//usage:
//usage:#define taskset_example_usage
//usage:       "$ taskset 0x7 ./dgemm_test&\n"
//usage:       "$ taskset -p 0x1 $!\n"
//usage:       "pid 4790's current affinity mask: 7\n"
//usage:       "pid 4790's new affinity mask: 1\n"
//usage:       "$ taskset 0x7 /bin/sh -c './taskset -p 0x1 $$'\n"
//usage:       "pid 6671's current affinity mask: 1\n"
//usage:       "pid 6671's new affinity mask: 1\n"
//usage:       "$ taskset -p 1\n"
//usage:       "pid 1's current affinity mask: 3\n"

#include <sched.h>
#include "libbb.h"

#if ENABLE_FEATURE_TASKSET_FANCY
#define TASKSET_PRINTF_MASK "%s"
/* craft a string from the mask */
static char *from_cpuset(cpu_set_t *mask)
{
	int i;
	char *ret = NULL;
	char *str = xzalloc((CPU_SETSIZE / 4) + 1); /* we will leak it */

	for (i = CPU_SETSIZE - 4; i >= 0; i -= 4) {
		int val = 0;
		int off;
		for (off = 0; off <= 3; ++off)
			if (CPU_ISSET(i + off, mask))
				val |= 1 << off;
		if (!ret && val)
			ret = str;
		*str++ = bb_hexdigits_upcase[val] | 0x20;
	}
	return ret;
}
#else
#define TASKSET_PRINTF_MASK "%llx"
static unsigned long long from_cpuset(cpu_set_t *mask)
{
	struct BUG_CPU_SETSIZE_is_too_small {
		char BUG_CPU_SETSIZE_is_too_small[
			CPU_SETSIZE < sizeof(int) ? -1 : 1];
	};
	char *p = (void*)mask;

	/* Take the least significant bits. Careful!
	 * Consider both CPU_SETSIZE=4 and CPU_SETSIZE=1024 cases
	 */
#if BB_BIG_ENDIAN
	/* For big endian, it means LAST bits */
	if (CPU_SETSIZE < sizeof(long))
		p += CPU_SETSIZE - sizeof(int);
	else if (CPU_SETSIZE < sizeof(long long))
		p += CPU_SETSIZE - sizeof(long);
	else
		p += CPU_SETSIZE - sizeof(long long);
#endif
	if (CPU_SETSIZE < sizeof(long))
		return *(unsigned*)p;
	if (CPU_SETSIZE < sizeof(long long))
		return *(unsigned long*)p;
	return *(unsigned long long*)p;
}
#endif


int taskset_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
int taskset_main(int argc UNUSED_PARAM, char **argv)
{
	cpu_set_t mask;
	pid_t pid = 0;
	unsigned opt_p;
	const char *current_new;
	char *pid_str;
	char *aff = aff; /* for compiler */

	/* NB: we mimic util-linux's taskset: -p does not take
	 * an argument, i.e., "-pN" is NOT valid, only "-p N"!
	 * Indeed, util-linux-2.13-pre7 uses:
	 * getopt_long(argc, argv, "+pchV", ...), not "...p:..." */

	opt_complementary = "-1"; /* at least 1 arg */
	opt_p = getopt32(argv, "+p");
	argv += optind;

	if (opt_p) {
		pid_str = *argv++;
		if (*argv) { /* "-p <aff> <pid> ...rest.is.ignored..." */
			aff = pid_str;
			pid_str = *argv; /* NB: *argv != NULL in this case */
		}
		/* else it was just "-p <pid>", and *argv == NULL */
		pid = xatoul_range(pid_str, 1, ((unsigned)(pid_t)ULONG_MAX) >> 1);
	} else {
		aff = *argv++; /* <aff> <cmd...> */
		if (!*argv)
			bb_show_usage();
	}

	current_new = "current\0new";
	if (opt_p) {
 print_aff:
		if (sched_getaffinity(pid, sizeof(mask), &mask) < 0)
			bb_perror_msg_and_die("can't %cet pid %d's affinity", 'g', pid);
		printf("pid %d's %s affinity mask: "TASKSET_PRINTF_MASK"\n",
				pid, current_new, from_cpuset(&mask));
		if (!*argv) {
			/* Either it was just "-p <pid>",
			 * or it was "-p <aff> <pid>" and we came here
			 * for the second time (see goto below) */
			return EXIT_SUCCESS;
		}
		*argv = NULL;
		current_new += 8; /* "new" */
	}

	{ /* Affinity was specified, translate it into cpu_set_t */
		unsigned i;
		/* Do not allow zero mask: */
		unsigned long long m = xstrtoull_range(aff, 0, 1, ULLONG_MAX);
		enum { CNT_BIT = CPU_SETSIZE < sizeof(m)*8 ? CPU_SETSIZE : sizeof(m)*8 };

		CPU_ZERO(&mask);
		for (i = 0; i < CNT_BIT; i++) {
			unsigned long long bit = (1ULL << i);
			if (bit & m)
				CPU_SET(i, &mask);
		}
	}

	/* Set pid's or our own (pid==0) affinity */
	if (sched_setaffinity(pid, sizeof(mask), &mask))
		bb_perror_msg_and_die("can't %cet pid %d's affinity", 's', pid);

	if (!argv[0]) /* "-p <aff> <pid> [...ignored...]" */
		goto print_aff; /* print new affinity and exit */

	BB_EXECVP_or_die(argv);
}