CPU affinity patch for 'nice'

[Roger Venable]

Hello!  I wrote a patch for 'nice' which allows setting CPU affinity in
kernel 2.6 on SMP machines.  This patch applies to the source code in
coreutils-5.2.1-7.src.rpm, file src/nice.c , and it compiles on a Fedora
Core 2 distribution with kernel 2.6.8-1.521smp.  Additional code may be
needed to ensure the availability of the scheduler functions at run-time,
but I am not quite sure as to the best way to do this.

The page at http://www.gnu.org/software/coreutils/ cites this email
address as the place to suggest this patch as an improvement.  I would be
interested in your opinions, thank you very much for your consideration.

Roger Cody Venable
Ann Arbor, Michigan


--- src/nice.c.original 2004-09-12 15:48:53.594978715 -0400
+++ src/nice.c  2004-09-12 17:27:50.463900716 -0400
@@ -20,6 +20,12 @@
 #include <config.h>
 #include <stdio.h>

+/*
+#include <stdlib.h>
+#include <sched.h>
+#include <linux/unistd.h>
+*/
+
 #include <assert.h>

 #include <getopt.h>
@@ -33,6 +39,19 @@
 # include <sys/resource.h>
 #endif

+/*
+ * provide the proper syscall information if our libc
+ * is not yet updated.
+ */
+/*
+#ifndef __NR_sched_setaffinity
+#define __NR_sched_setaffinity241
+#define __NR_sched_getaffinity242
+_syscall3 (int, sched_setaffinity, pid_t, pid, unsigned int, len,
unsigned long *, user_mask_ptr)
+_syscall3 (int, sched_getaffinity, pid_t, pid, unsigned int, len,
unsigned long *, user_mask_ptr)
+#endif
+*/
+
 #include "error.h"
 #include "long-options.h"
 #include "posixver.h"
@@ -55,6 +74,7 @@
 static struct option const longopts[] =
 {
   {"adjustment", required_argument, NULL, 'n'},
+  {"affinity", required_argument, NULL, 'c'},
   {NULL, 0, NULL, 0}
 };

@@ -88,7 +108,13 @@
   long int adjustment = 0;
   int minusflag = 0;
   int adjustment_given = 0;
+  int affinity_given = 0;
   int i;
+  unsigned long new_mask = 0;
+  unsigned int len_mask = sizeof(new_mask);
+  unsigned long cur_mask = 0;
+  pid_t p = 0;
+  int ret;

   initialize_main (&argc, &argv);
   program_name = argv[0];
@@ -137,7 +163,7 @@
          /* Initialize getopt_long's internal state.  */
          optind = 0;

-         if ((optc = getopt_long (argc - (i - 1), fake_argv, "+n:",
+         if ((optc = getopt_long (argc - (i - 1), fake_argv, "+n:c:",
                                   longopts, NULL)) != -1)
            {
              switch (optc)
@@ -153,6 +179,17 @@
                  minusflag = 0;
                  adjustment_given = 1;
                  break;
+
+                case 'c':
+                  if (xstrtol (optarg, NULL, 10, &new_mask, "")
+                      != LONGINT_OK)
+                    error (EXIT_FAIL, 0, _("invalid affinity `%s'"),
optarg);
+
+                  if (new_mask <= 0)
+                    error (EXIT_FAIL, 0, _("invalid affinity `%s'"),
optarg);
+
+                  affinity_given = 1;
+                  break;
                }
            }

@@ -170,9 +207,9 @@

   if (i == argc)
     {
-      if (adjustment_given)
+      if (adjustment_given || affinity_given)
        {
-         error (0, 0, _("a command must be given with an adjustment"));
+         error (0, 0, _("a command must be given with an adjustment or
affinity"));
          usage (EXIT_FAIL);
        }
       /* No command given; print the priority. */
@@ -195,6 +232,18 @@
 #endif
     error (EXIT_FAIL, errno, _("cannot set priority"));

+  if ( affinity_given )
+  {
+    ret = sched_getaffinity(p, len_mask, &cur_mask);
+    if (( ret == 0 ) && ( new_mask != cur_mask ))
+    {
+      new_mask = new_mask & cur_mask; /* limit to available CPUs */
+      ret = sched_setaffinity(p, len_mask, &new_mask);
+      if ( ret != 0 )
+        error (EXIT_FAIL, errno, _("cannot set affinity"));
+    }
+  }
+
   execvp (argv[i], &argv[i]);

   {