3 files changed, 65 insertions, 2 deletions
diff --git a/src/libutil/Makefile.am b/src/libutil/Makefile.am
index fe896eec5012..0c4073e66eab 100644
--- a/src/libutil/Makefile.am
+++ b/src/libutil/Makefile.am
@@ -1,12 +1,12 @@
 pkglib_LTLIBRARIES = libutil.la
 
 libutil_la_SOURCES = util.cc hash.cc serialise.cc \
-  archive.cc xml-writer.cc
+  archive.cc xml-writer.cc affinity.cc
 
 libutil_la_LIBADD = ../boost/format/libformat.la
 
 pkginclude_HEADERS = util.hh hash.hh serialise.hh \
-  archive.hh xml-writer.hh types.hh
+  archive.hh xml-writer.hh types.hh affinity.hh
 
 if !HAVE_OPENSSL
 libutil_la_SOURCES += \
diff --git a/src/libutil/affinity.cc b/src/libutil/affinity.cc
new file mode 100644
index 000000000000..3a20fd2774fe
--- /dev/null
+++ b/src/libutil/affinity.cc
@@ -0,0 +1,54 @@
+#include "types.hh"
+#include "util.hh"
+#include "affinity.hh"
+
+#if HAVE_SCHED_H
+#include <sched.h>
+#endif
+
+namespace nix {
+
+
+static bool didSaveAffinity = false;
+static cpu_set_t savedAffinity;
+
+
+void setAffinityTo(int cpu)
+{
+#if HAVE_SCHED_SETAFFINITY
+    if (sched_getaffinity(0, sizeof(cpu_set_t), &savedAffinity) == -1) return;
+    didSaveAffinity = true;
+    printMsg(lvlDebug, format("locking this thread to CPU %1%") % cpu);
+    cpu_set_t newAffinity;
+    CPU_ZERO(&newAffinity);
+    CPU_SET(cpu, &newAffinity);
+    if (sched_setaffinity(0, sizeof(cpu_set_t), &newAffinity) == -1)
+        printMsg(lvlError, format("failed to lock thread to CPU %1%") % cpu);
+#endif
+}
+
+
+int lockToCurrentCPU()
+{
+#if HAVE_SCHED_SETAFFINITY
+    if (getEnv("NIX_AFFINITY_HACK", "1") == "1") {
+        int cpu = sched_getcpu();
+        if (cpu != -1) setAffinityTo(cpu);
+        return cpu;
+    }
+#endif
+    return -1;
+}
+
+
+void restoreAffinity()
+{
+#if HAVE_SCHED_SETAFFINITY
+    if (!didSaveAffinity) return;
+    if (sched_setaffinity(0, sizeof(cpu_set_t), &savedAffinity) == -1)
+        printMsg(lvlError, "failed to restore affinity %1%");
+#endif
+}
+
+
+}
diff --git a/src/libutil/affinity.hh b/src/libutil/affinity.hh
new file mode 100644
index 000000000000..c1bd28e1367a
--- /dev/null
+++ b/src/libutil/affinity.hh
@@ -0,0 +1,9 @@
+#pragma once
+
+namespace nix {
+
+void setAffinityTo(int cpu);
+int lockToCurrentCPU();
+void restoreAffinity();
+
+}