tools/tests/affinity/tags.c

   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
  15 int             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
  31 static void
  32 usage()
  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
  43 void
  44 thread_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
  60 int
  61 thread_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
  79 char                    input[81];
  80 int
  81 main(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 }