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