From 91539d305ff035d53a6de75f8af1ebbd7df4e622 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <eelco.dolstra@logicblox.com>
Date: Fri, 22 Apr 2016 18:19:48 +0200
Subject: [PATCH] nix copy: Parallelise

---
 src/nix/copy.cc         | 75 ++++++++++++++++++++++++++++++++---------
 src/nix/progress-bar.cc |  2 +-
 src/nix/sigs.cc         |  3 --
 src/nix/verify.cc       |  3 --
 4 files changed, 61 insertions(+), 22 deletions(-)

diff --git a/src/nix/copy.cc b/src/nix/copy.cc
index 87e6f604a..16b16910c 100644
--- a/src/nix/copy.cc
+++ b/src/nix/copy.cc
@@ -3,6 +3,7 @@
 #include "shared.hh"
 #include "store-api.hh"
 #include "sync.hh"
+#include "thread-pool.hh"
 
 #include <atomic>
 
@@ -57,29 +58,73 @@ struct CmdCopy : StorePathsCommand
 
         progressBar.updateStatus(showProgress());
 
-        storePaths.reverse(); // FIXME: assumes reverse topo sort
+        struct Graph
+        {
+            std::set<Path> left;
+            std::map<Path, std::set<Path>> refs, rrefs;
+        };
 
-        for (auto & storePath : storePaths) {
+        Sync<Graph> graph_;
+        {
+            auto graph(graph_.lock());
+            graph->left = PathSet(storePaths.begin(), storePaths.end());
+        }
+
+        ThreadPool pool;
+
+        std::function<void(const Path &)> doPath;
+
+        doPath = [&](const Path & storePath) {
             checkInterrupt();
 
-            if (dstStore->isValidPath(storePath)) {
+            if (!dstStore->isValidPath(storePath)) {
+                auto activity(progressBar.startActivity(format("copying ‘%s’...") % storePath));
+
+                StringSink sink;
+                srcStore->exportPaths({storePath}, false, sink);
+
+                StringSource source(*sink.s);
+                dstStore->importPaths(false, source, 0);
+
+                done++;
+            } else
                 total--;
-                progressBar.updateStatus(showProgress());
-                continue;
-            }
 
-            auto activity(progressBar.startActivity(format("copying ‘%s’...") % storePath));
-
-            StringSink sink;
-            srcStore->exportPaths({storePath}, false, sink);
-
-            StringSource source(*sink.s);
-            dstStore->importPaths(false, source, 0);
-
-            done++;
             progressBar.updateStatus(showProgress());
+
+            /* Enqueue all paths that were waiting for this one. */
+            {
+                auto graph(graph_.lock());
+                graph->left.erase(storePath);
+                for (auto & rref : graph->rrefs[storePath]) {
+                    auto & refs(graph->refs[rref]);
+                    auto i = refs.find(storePath);
+                    assert(i != refs.end());
+                    refs.erase(i);
+                    if (refs.empty())
+                        pool.enqueue(std::bind(doPath, rref));
+                }
+            }
+        };
+
+        /* Build the dependency graph; enqueue all paths with no
+           dependencies. */
+        for (auto & storePath : storePaths) {
+            auto info = srcStore->queryPathInfo(storePath);
+            {
+                auto graph(graph_.lock());
+                for (auto & ref : info->references)
+                    if (ref != storePath && graph->left.count(ref)) {
+                        graph->refs[storePath].insert(ref);
+                        graph->rrefs[ref].insert(storePath);
+                    }
+                if (graph->refs[storePath].empty())
+                    pool.enqueue(std::bind(doPath, storePath));
+            }
         }
 
+        pool.process();
+
         progressBar.done();
     }
 };
diff --git a/src/nix/progress-bar.cc b/src/nix/progress-bar.cc
index ed7b578e2..746b891a7 100644
--- a/src/nix/progress-bar.cc
+++ b/src/nix/progress-bar.cc
@@ -29,12 +29,12 @@ void ProgressBar::updateStatus(const std::string & s)
 
 void ProgressBar::done()
 {
+    _writeToStderr = decltype(_writeToStderr)();
     auto state_(state.lock());
     assert(state_->activities.empty());
     state_->done = true;
     std::cerr << "\r\e[K";
     std::cerr.flush();
-    _writeToStderr = decltype(_writeToStderr)();
 }
 
 void ProgressBar::render(State & state_)
diff --git a/src/nix/sigs.cc b/src/nix/sigs.cc
index 56f7d2213..69073d884 100644
--- a/src/nix/sigs.cc
+++ b/src/nix/sigs.cc
@@ -1,4 +1,3 @@
-#include "affinity.hh" // FIXME
 #include "command.hh"
 #include "progress-bar.hh"
 #include "shared.hh"
@@ -31,8 +30,6 @@ struct CmdCopySigs : StorePathsCommand
 
     void run(ref<Store> store, Paths storePaths) override
     {
-        restoreAffinity(); // FIXME
-
         if (substituterUris.empty())
             throw UsageError("you must specify at least one substituter using ‘-s’");
 
diff --git a/src/nix/verify.cc b/src/nix/verify.cc
index 3844535e7..20ce26fa7 100644
--- a/src/nix/verify.cc
+++ b/src/nix/verify.cc
@@ -1,4 +1,3 @@
-#include "affinity.hh" // FIXME
 #include "command.hh"
 #include "progress-bar.hh"
 #include "shared.hh"
@@ -52,8 +51,6 @@ struct CmdVerify : StorePathsCommand
 
     void run(ref<Store> store, Paths storePaths) override
     {
-        restoreAffinity(); // FIXME
-
         std::vector<ref<Store>> substituters;
         for (auto & s : substituterUris)
             substituters.push_back(openStoreAt(s));