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