nix/libutil/affinity.cc

   1 #include "types.hh"
   2 #include "util.hh"
   3 #include "affinity.hh"
   4
   5 #if HAVE_SCHED_H
   6 #include <sched.h>
   7 #endif
   8
   9 namespace nix {
  10
  11
  12 #if HAVE_SCHED_SETAFFINITY
  13 static bool didSaveAffinity = false;
  14 static cpu_set_t savedAffinity;
  15 #endif
  16
  17
  18 void setAffinityTo(int cpu)
  19 {
  20 #if HAVE_SCHED_SETAFFINITY
  21     if (sched_getaffinity(0, sizeof(cpu_set_t), &savedAffinity) == -1) return;
  22     didSaveAffinity = true;
  23     printMsg(lvlDebug, format("locking this thread to CPU %1%") % cpu);
  24     cpu_set_t newAffinity;
  25     CPU_ZERO(&newAffinity);
  26     CPU_SET(cpu, &newAffinity);
  27     if (sched_setaffinity(0, sizeof(cpu_set_t), &newAffinity) == -1)
  28         printMsg(lvlError, format("failed to lock thread to CPU %1%") % cpu);
  29 #endif
  30 }
  31
  32
  33 int lockToCurrentCPU()
  34 {
  35 #if HAVE_SCHED_SETAFFINITY
  36     int cpu = sched_getcpu();
  37     if (cpu != -1) setAffinityTo(cpu);
  38     return cpu;
  39 #else
  40     return -1;
  41 #endif
  42 }
  43
  44
  45 void restoreAffinity()
  46 {
  47 #if HAVE_SCHED_SETAFFINITY
  48     if (!didSaveAffinity) return;
  49     if (sched_setaffinity(0, sizeof(cpu_set_t), &savedAffinity) == -1)
  50         printMsg(lvlError, "failed to restore affinity %1%");
  51 #endif
  52 }
  53
  54
  55 }