| 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 | } |