Factor out parallel processing of work items that have dependencies

2 files changed, 78 insertions, 55 deletions

diff --git a/src/libutil/thread-pool.hh b/src/libutil/thread-pool.hh
index 939bcf1ef93b..78b63467d62e 100644
--- a/src/libutil/thread-pool.hh
+++ b/src/libutil/thread-pool.hh
@@ -6,6 +6,7 @@
 #include <queue>
 #include <functional>
 #include <thread>
+#include <map>
 
 namespace nix {
 
@@ -54,4 +55,63 @@ private:
     void workerEntry();
 };
 
+/* Process in parallel a set of items of type T that have a partial
+   ordering between them. Thus, any item is only processed after all
+   its dependencies have been processed. */
+template<typename T>
+void processGraph(
+    ThreadPool & pool,
+    const std::set<T> & nodes,
+    std::function<std::set<T>(const T &)> getEdges,
+    std::function<void(const T &)> processNode)
+{
+    struct Graph {
+        std::set<T> left;
+        std::map<T, std::set<T>> refs, rrefs;
+        std::function<void(T)> wrap;
+    };
+
+    ref<Sync<Graph>> graph_ = make_ref<Sync<Graph>>();
+
+    auto wrapWork = [&pool, graph_, processNode](const T & node) {
+        processNode(node);
+
+        /* Enqueue work for all nodes that were waiting on this one. */
+        {
+            auto graph(graph_->lock());
+            graph->left.erase(node);
+            for (auto & rref : graph->rrefs[node]) {
+                auto & refs(graph->refs[rref]);
+                auto i = refs.find(node);
+                assert(i != refs.end());
+                refs.erase(i);
+                if (refs.empty())
+                    pool.enqueue(std::bind(graph->wrap, rref));
+            }
+        }
+    };
+
+    {
+        auto graph(graph_->lock());
+        graph->left = nodes;
+        graph->wrap = wrapWork;
+    }
+
+    /* Build the dependency graph; enqueue all nodes with no
+       dependencies. */
+    for (auto & node : nodes) {
+        auto refs = getEdges(node);
+        {
+            auto graph(graph_->lock());
+            for (auto & ref : refs)
+                if (ref != node && graph->left.count(ref)) {
+                    graph->refs[node].insert(ref);
+                    graph->rrefs[ref].insert(node);
+                }
+            if (graph->refs[node].empty())
+                pool.enqueue(std::bind(graph->wrap, node));
+        }
+    }
+}
+
 }
diff --git a/src/nix/copy.cc b/src/nix/copy.cc
index 16b16910c46e..b5bd362d63a6 100644
--- a/src/nix/copy.cc
+++ b/src/nix/copy.cc
@@ -58,70 +58,33 @@ struct CmdCopy : StorePathsCommand
 
         progressBar.updateStatus(showProgress());
 
-        struct Graph
-        {
-            std::set<Path> left;
-            std::map<Path, std::set<Path>> refs, rrefs;
-        };
-
-        Sync<Graph> graph_;
-        {
-            auto graph(graph_.lock());
-            graph->left = PathSet(storePaths.begin(), storePaths.end());
-        }
-
         ThreadPool pool;
 
-        std::function<void(const Path &)> doPath;
+        processGraph<Path>(pool,
+            PathSet(storePaths.begin(), storePaths.end()),
 
-        doPath = [&](const Path & storePath) {
-            checkInterrupt();
-
-            if (!dstStore->isValidPath(storePath)) {
-                auto activity(progressBar.startActivity(format("copying ‘%s’...") % storePath));
+            [&](const Path & storePath) {
+                return srcStore->queryPathInfo(storePath)->references;
+            },
 
-                StringSink sink;
-                srcStore->exportPaths({storePath}, false, sink);
+            [&](const Path & storePath) {
+                checkInterrupt();
 
-                StringSource source(*sink.s);
-                dstStore->importPaths(false, source, 0);
+                if (!dstStore->isValidPath(storePath)) {
+                    auto activity(progressBar.startActivity(format("copying ‘%s’...") % storePath));
 
-                done++;
-            } else
-                total--;
+                    StringSink sink;
+                    srcStore->exportPaths({storePath}, false, sink);
 
-            progressBar.updateStatus(showProgress());
+                    StringSource source(*sink.s);
+                    dstStore->importPaths(false, source, 0);
 
-            /* 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));
-                }
-            }
-        };
+                    done++;
+                } else
+                    total--;
 
-        /* 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));
-            }
-        }
+                progressBar.updateStatus(showProgress());
+            });
 
         pool.process();