8 files changed, 88 insertions, 4 deletions
diff --git a/src/libstore/build.cc b/src/libstore/build.cc
index 1ca99dcc5cd4..67e38da0bb0f 100644
--- a/src/libstore/build.cc
+++ b/src/libstore/build.cc
@@ -7,6 +7,7 @@
 #include "local-store.hh"
 #include "util.hh"
 #include "archive.hh"
+#include "affinity.hh"
 
 #include <map>
 #include <sstream>
@@ -366,6 +367,8 @@ void Goal::trace(const format & f)
 /* Common initialisation performed in child processes. */
 static void commonChildInit(Pipe & logPipe)
 {
+    restoreAffinity();
+
     /* Put the child in a separate session (and thus a separate
        process group) so that it has no controlling terminal (meaning
        that e.g. ssh cannot open /dev/tty) and it doesn't receive
@@ -568,6 +571,7 @@ static void runSetuidHelper(const string & command,
             args.push_back(0);
 
             restoreSIGPIPE();
+            restoreAffinity();
 
             execve(program.c_str(), (char * *) &args[0], 0);
             throw SysError(format("executing `%1%'") % program);
diff --git a/src/libstore/local-store.cc b/src/libstore/local-store.cc
index 29095e1eaa0f..34b4a5158d94 100644
--- a/src/libstore/local-store.cc
+++ b/src/libstore/local-store.cc
@@ -5,6 +5,7 @@
 #include "pathlocks.hh"
 #include "worker-protocol.hh"
 #include "derivations.hh"
+#include "affinity.hh"
 
 #include <iostream>
 #include <algorithm>
@@ -1021,6 +1022,7 @@ void LocalStore::startSubstituter(const Path & substituter, RunningSubstituter &
 
     case 0: /* child */
         try {
+            restoreAffinity();
             if (dup2(toPipe.readSide, STDIN_FILENO) == -1)
                 throw SysError("dupping stdin");
             if (dup2(fromPipe.writeSide, STDOUT_FILENO) == -1)
diff --git a/src/libstore/remote-store.cc b/src/libstore/remote-store.cc
index 2b5a932131c3..3764b4813a4e 100644
--- a/src/libstore/remote-store.cc
+++ b/src/libstore/remote-store.cc
@@ -3,6 +3,7 @@
 #include "remote-store.hh"
 #include "worker-protocol.hh"
 #include "archive.hh"
+#include "affinity.hh"
 #include "globals.hh"
 
 #include <sys/types.h>
@@ -15,7 +16,6 @@
 #include <unistd.h>
 #include <cstring>
 
-
 namespace nix {
 
 
@@ -71,8 +71,19 @@ void RemoteStore::openConnection(bool reserveSpace)
         if (GET_PROTOCOL_MAJOR(daemonVersion) != GET_PROTOCOL_MAJOR(PROTOCOL_VERSION))
             throw Error("Nix daemon protocol version not supported");
         writeInt(PROTOCOL_VERSION, to);
+
+        if (GET_PROTOCOL_MINOR(daemonVersion) >= 14) {
+            int cpu = lockToCurrentCPU();
+            if (cpu != -1) {
+                writeInt(1, to);
+                writeInt(cpu, to);
+            } else
+                writeInt(0, to);
+        }
+
         if (GET_PROTOCOL_MINOR(daemonVersion) >= 11)
             writeInt(reserveSpace, to);
+
         processStderr();
     }
     catch (Error & e) {
diff --git a/src/libstore/worker-protocol.hh b/src/libstore/worker-protocol.hh
index 07f825b9273f..9317f89c3763 100644
--- a/src/libstore/worker-protocol.hh
+++ b/src/libstore/worker-protocol.hh
@@ -6,7 +6,7 @@ namespace nix {
 #define WORKER_MAGIC_1 0x6e697863
 #define WORKER_MAGIC_2 0x6478696f
 
-#define PROTOCOL_VERSION 0x10d
+#define PROTOCOL_VERSION 0x10e
 #define GET_PROTOCOL_MAJOR(x) ((x) & 0xff00)
 #define GET_PROTOCOL_MINOR(x) ((x) & 0x00ff)
 
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();
+
+}
diff --git a/src/nix-daemon/nix-daemon.cc b/src/nix-daemon/nix-daemon.cc
index a6f29b7c9795..86f5c0a244eb 100644
--- a/src/nix-daemon/nix-daemon.cc
+++ b/src/nix-daemon/nix-daemon.cc
@@ -4,6 +4,7 @@
 #include "serialise.hh"
 #include "worker-protocol.hh"
 #include "archive.hh"
+#include "affinity.hh"
 #include "globals.hh"
 
 #include <cstring>
@@ -671,6 +672,9 @@ static void processConnection(bool trusted)
     to.flush();
     unsigned int clientVersion = readInt(from);
 
+    if (GET_PROTOCOL_MINOR(clientVersion) >= 14 && readInt(from))
+        setAffinityTo(readInt(from));
+
     bool reserveSpace = true;
     if (GET_PROTOCOL_MINOR(clientVersion) >= 11)
         reserveSpace = readInt(from) != 0;