1#include <AvailabilityMacros.h>
2#ifdef AVAILABLE_MAC_OS_X_VERSION_10_5_AND_LATER
3#include </System/Library/Frameworks/System.framework/PrivateHeaders/mach/thread_policy.h>
4#endif
5#include <mach/mach.h>
6#include <mach/mach_traps.h>
7#include <mach/mach_error.h>
8#include <mach/mach_time.h>
9#include <stdio.h>
10#include <stdlib.h>
11#include <string.h>
12#include <unistd.h>
13#include <err.h>
14
15int		verbosity = 1;
16
17#define DBG(x...) do {				\
18	if (verbosity > 1) {			\
19		printf(x);			\
20	}					\
21} while (0)
22
23#define mutter(x...) do {			\
24	if (verbosity > 0) {			\
25		printf(x);			\
26	}					\
27} while (0)
28
29#define s_if_plural(x)	(((x) > 1) ? "s" : "")
30
31static void
32usage()
33{
34	fprintf(stderr,
35		"usage: tags [-i]   interactive/input\n"
36		"            [-v V] verbosity level 0..2 (1)\n"
37		"            [-h]   help info\n"
38		"            pid    process id of target task\n"
39	);
40	exit(1);
41}
42
43void
44thread_tag_set(thread_t thread, int tag)
45{
46	kern_return_t			ret;
47	thread_affinity_policy_data_t	policy;
48
49	policy.affinity_tag = tag;
50	ret = thread_policy_set(
51			thread, THREAD_AFFINITY_POLICY,
52			(thread_policy_t) &policy,
53			THREAD_AFFINITY_POLICY_COUNT);
54	if (ret != KERN_SUCCESS) {
55		printf("thread_policy_set(1) returned %d\n", ret);
56		exit(1);
57	}
58}
59
60int
61thread_tag_get(thread_t thread)
62{
63	kern_return_t			ret;
64	boolean_t			get_default = FALSE;
65	thread_affinity_policy_data_t	policy;
66	mach_msg_type_number_t		count = THREAD_AFFINITY_POLICY_COUNT;
67
68	ret = thread_policy_get(
69			thread, THREAD_AFFINITY_POLICY,
70			(thread_policy_t) &policy, &count, &get_default);
71	if (ret != KERN_SUCCESS) {
72		printf("thread_policy_set(1) returned %d\n", ret);
73		exit(1);
74	}
75
76	return policy.affinity_tag;
77}
78
79char			input[81];
80int
81main(int argc, char *argv[])
82{
83	kern_return_t		ret;
84	mach_port_name_t	port;
85	int			pid;
86	int			c;
87	thread_act_t		*thread_array;
88	mach_msg_type_number_t	num_threads;
89	int			i;
90	boolean_t		interactive = FALSE;
91	int			tag;
92
93	if (geteuid() != 0) {
94		printf("Must be run as root\n");
95		exit(1);
96	}
97
98	/* Do switch parsing: */
99	while ((c = getopt (argc, argv, "hiv:")) != -1) {
100		switch (c) {
101		case 'i':
102			interactive = TRUE;
103			break;
104		case 'v':
105			verbosity = atoi(optarg);
106			break;
107		case 'h':
108		case '?':
109		default:
110			usage();
111		}
112	}
113	argc -= optind; argv += optind;
114	if (argc > 0)
115		pid = atoi(*argv);
116
117	ret = task_for_pid(mach_task_self(), pid, &port);
118	if (ret != KERN_SUCCESS)
119		err(1, "task_for_pid(,%d,) returned %d", pid, ret);
120
121	mutter("task %p\n", port);
122	ret = task_threads(port, &thread_array, &num_threads);
123	if (ret != KERN_SUCCESS)
124		err(1, "task_threads() returned %d", pid, ret);
125
126	for (i = 0; i < num_threads; i++) {
127		printf(" %d: thread 0x%08x tag %d\n",
128			i, thread_array[i], thread_tag_get(thread_array[i]));
129	}
130
131	while (interactive) {
132		printf("Enter new tag or <return> to skip or ^D to quit\n");
133		for (i = 0; i < num_threads; i++) {
134			tag =  thread_tag_get(thread_array[i]);
135			printf(" %d: thread 0x%08x tag %d: ",
136				i, thread_array[i], tag);
137			fflush(stdout);
138			(void) fgets(input, 20, stdin);
139			if (feof(stdin)) {
140				printf("\n");
141				interactive = FALSE;
142				break;
143			}
144			if (strlen(input) > 1) {
145				tag = atoi(input);
146				thread_tag_set(thread_array[i], tag);
147			}
148		}
149	}
150
151	return 0;
152}
153