src/psched.c

   1 /*
   2  * netsniff-ng - the packet sniffing beast
   3  * By Daniel Borkmann <daniel@netsniff-ng.org>
   4  * Copyright 2009-2011 Daniel Borkmann.
   5  * Copyright 2010 Emmanuel Roullit.
   6  * Subject to the GPL.
   7  */
   8
   9 /* Process RT scheduling */
  10
  11 #define _GNU_SOURCE
  12 #include <stdio.h>
  13 #include <stdlib.h>
  14 #include <string.h>
  15 #include <unistd.h>
  16 #include <errno.h>
  17
  18 #include "psched.h"
  19 #include "die.h"
  20
  21 static inline const char *next_token(const char *q, int sep)
  22 {
  23         if (q)
  24                 q = strchr(q, sep);
  25         if (q)
  26                 q++;
  27
  28         return (q);
  29 }
  30
  31 int set_cpu_affinity(const char *str, int inverted)
  32 {
  33         int ret, i, cpus;
  34         const char *p, *q;
  35         cpu_set_t cpu_bitmask;
  36
  37         q = str;
  38         cpus = sysconf(_SC_NPROCESSORS_CONF);
  39         CPU_ZERO(&cpu_bitmask);
  40
  41         for (i = 0; inverted && i < cpus; ++i)
  42                 CPU_SET(i, &cpu_bitmask);
  43
  44         while (p = q, q = next_token(q, ','), p) {
  45                 unsigned int a;  /* Beginning of range */
  46                 unsigned int b;  /* End of range */
  47                 unsigned int s;  /* Stride */
  48                 const char *c1, *c2;
  49
  50                 if (sscanf(p, "%u", &a) < 1)
  51                         return -EINVAL;
  52
  53                 b = a;
  54                 s = 1;
  55
  56                 c1 = next_token(p, '-');
  57                 c2 = next_token(p, ',');
  58
  59                 if (c1 != NULL && (c2 == NULL || c1 < c2)) {
  60                         if (sscanf(c1, "%u", &b) < 1)
  61                                 return -EINVAL;
  62                         c1 = next_token(c1, ':');
  63                         if (c1 != NULL && (c2 == NULL || c1 < c2))
  64                                 if (sscanf(c1, "%u", &s) < 1)
  65                                         return -EINVAL;
  66                 }
  67
  68                 if (!(a <= b))
  69                         return -EINVAL;
  70
  71                 while (a <= b) {
  72                         if (inverted)
  73                                 CPU_CLR(a, &cpu_bitmask);
  74                         else
  75                                 CPU_SET(a, &cpu_bitmask);
  76                         a += s;
  77                 }
  78         }
  79
  80         ret = sched_setaffinity(getpid(), sizeof(cpu_bitmask),
  81                                 &cpu_bitmask);
  82         if (ret)
  83                 panic("Can't set this cpu affinity!\n");
  84         return 0;
  85 }
  86
  87 char *get_cpu_affinity(char *cpu_string, size_t len)
  88 {
  89         int ret, i, cpu;
  90         cpu_set_t cpu_bitmask;
  91
  92         if (len != sysconf(_SC_NPROCESSORS_CONF) + 1)
  93                 return NULL;
  94         CPU_ZERO(&cpu_bitmask);
  95
  96         ret = sched_getaffinity(getpid(), sizeof(cpu_bitmask),
  97                                 &cpu_bitmask);
  98         if (ret) {
  99                 whine("Can't fetch cpu affinity!\n");
 100                 return NULL;
 101         }
 102
 103         for (i = 0, cpu_string[len - 1] = 0; i < len - 1; ++i) {
 104                 cpu = CPU_ISSET(i, &cpu_bitmask);
 105                 cpu_string[i] = (cpu ? '1' : '0');
 106         }
 107
 108         return cpu_string;
 109 }
 110
 111 int set_proc_prio(int priority)
 112 {
 113         /*
 114          * setpriority() is clever, even if you put a nice value which
 115          * is out of range it corrects it to the closest valid nice value
 116          */
 117         int ret = setpriority(PRIO_PROCESS, getpid(), priority);
 118         if (ret)
 119                 panic("Can't set nice val to %i!\n", priority);
 120         return 0;
 121 }
 122
 123 int set_sched_status(int policy, int priority)
 124 {
 125         int ret, min_prio, max_prio;
 126         struct sched_param sp;
 127
 128         max_prio = sched_get_priority_max(policy);
 129         min_prio = sched_get_priority_min(policy);
 130
 131         if (max_prio == -1 || min_prio == -1)
 132                 whine("Cannot determine scheduler prio limits!\n");
 133         else if (priority < min_prio)
 134                 priority = min_prio;
 135         else if (priority > max_prio)
 136                 priority = max_prio;
 137
 138         memset(&sp, 0, sizeof(sp));
 139         sp.sched_priority = priority;
 140
 141         ret = sched_setscheduler(getpid(), policy, &sp);
 142         if (ret) {
 143                 whine("Cannot set scheduler policy!\n");
 144                 return -EINVAL;
 145         }
 146
 147         ret = sched_setparam(getpid(), &sp);
 148         if (ret) {
 149                 whine("Cannot set scheduler prio!\n");
 150                 return -EINVAL;
 151         }
 152
 153         return 0;
 154 }
 155