From 3d9de41a5b86520ff0c994dc7ac1023d2a441083 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <edolstra@gmail.com>
Date: Tue, 29 Jun 2021 21:17:48 +0200
Subject: [PATCH] Hacky fast closure copying mechanism

---
 src/libcmd/installables.cc      | 29 +++++++++++++++++++++++++++++
 src/libstore/daemon.cc          | 11 +++++++++++
 src/libstore/remote-store.cc    | 10 ++++++++++
 src/libstore/remote-store.hh    |  2 ++
 src/libstore/store-api.hh       |  3 ++-
 src/libstore/worker-protocol.hh |  1 +
 6 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/src/libcmd/installables.cc b/src/libcmd/installables.cc
index 66ec63da2..ca4d509c2 100644
--- a/src/libcmd/installables.cc
+++ b/src/libcmd/installables.cc
@@ -12,6 +12,7 @@
 #include "eval-cache.hh"
 #include "url.hh"
 #include "registry.hh"
+#include "remote-store.hh"
 
 #include <regex>
 #include <queue>
@@ -378,6 +379,7 @@ DerivedPaths InstallableValue::toDerivedPaths()
     DerivedPaths res;
 
     std::map<StorePath, std::set<std::string>> drvsToOutputs;
+    RealisedPath::Set drvsToCopy;
 
     // Group by derivation, helps with .all in particular
     for (auto & drv : toDerivations()) {
@@ -385,11 +387,38 @@ DerivedPaths InstallableValue::toDerivedPaths()
         if (outputName == "")
             throw Error("derivation '%s' lacks an 'outputName' attribute", state->store->printStorePath(drv.drvPath));
         drvsToOutputs[drv.drvPath].insert(outputName);
+        drvsToCopy.insert(drv.drvPath);
     }
 
     for (auto & i : drvsToOutputs)
         res.push_back(DerivedPath::Built { i.first, i.second });
 
+    // FIXME: Temporary hack
+    if (state->store != state->buildStore) {
+        RealisedPath::Set closure;
+        RealisedPath::closure(*state->store, drvsToCopy, closure);
+
+        if (dynamic_cast<RemoteStore *>(&*state->buildStore)) {
+            StorePathSet closure2;
+            for (auto & p : closure)
+                closure2.insert(p.path());
+
+            auto valid = state->buildStore->queryValidPaths(closure2);
+            StorePathSet missing;
+            for (auto & p : closure2)
+                if (!valid.count(p)) missing.insert(p);
+
+            if (!missing.empty()) {
+                auto source = sinkToSource([&](Sink & sink) {
+                    state->store->exportPaths(missing, sink);
+                });
+                state->buildStore->importPaths(*source, NoCheckSigs);
+            }
+
+        } else
+            copyPaths(state->store, state->buildStore, closure);
+    }
+
     return res;
 }
 
diff --git a/src/libstore/daemon.cc b/src/libstore/daemon.cc
index e06fb9ce2..468152f7f 100644
--- a/src/libstore/daemon.cc
+++ b/src/libstore/daemon.cc
@@ -505,6 +505,17 @@ static void performOp(TunnelLogger * logger, ref<Store> store,
         break;
     }
 
+    case wopImportPaths2: {
+        logger->startWork();
+        auto paths = store->importPaths(from,
+            trusted ? NoCheckSigs : CheckSigs);
+        logger->stopWork();
+        Strings paths2;
+        for (auto & i : paths) paths2.push_back(store->printStorePath(i));
+        to << paths2;
+        break;
+    }
+
     case wopBuildPaths: {
         auto drvs = readDerivedPaths(*store, clientVersion, from);
         BuildMode mode = bmNormal;
diff --git a/src/libstore/remote-store.cc b/src/libstore/remote-store.cc
index aec243637..cd2d718ec 100644
--- a/src/libstore/remote-store.cc
+++ b/src/libstore/remote-store.cc
@@ -875,6 +875,16 @@ void RemoteStore::queryMissing(const std::vector<DerivedPath> & targets,
 }
 
 
+StorePaths RemoteStore::importPaths(Source & source, CheckSigsFlag checkSigs)
+{
+    auto conn(getConnection());
+    conn->to << wopImportPaths2;
+    source.drainInto(conn->to);
+    conn.processStderr();
+    return worker_proto::read(*this, conn->from, Phantom<StorePaths> {});
+}
+
+
 void RemoteStore::connect()
 {
     auto conn(getConnection());
diff --git a/src/libstore/remote-store.hh b/src/libstore/remote-store.hh
index 6cf76a46d..293393c32 100644
--- a/src/libstore/remote-store.hh
+++ b/src/libstore/remote-store.hh
@@ -112,6 +112,8 @@ public:
         StorePathSet & willBuild, StorePathSet & willSubstitute, StorePathSet & unknown,
         uint64_t & downloadSize, uint64_t & narSize) override;
 
+    StorePaths importPaths(Source & source, CheckSigsFlag checkSigs) override;
+
     void connect() override;
 
     unsigned int getProtocol() override;
diff --git a/src/libstore/store-api.hh b/src/libstore/store-api.hh
index 4faa65228..356ae615c 100644
--- a/src/libstore/store-api.hh
+++ b/src/libstore/store-api.hh
@@ -676,7 +676,7 @@ public:
        the Nix store. Optionally, the contents of the NARs are
        preloaded into the specified FS accessor to speed up subsequent
        access. */
-    StorePaths importPaths(Source & source, CheckSigsFlag checkSigs = CheckSigs);
+    virtual StorePaths importPaths(Source & source, CheckSigsFlag checkSigs = CheckSigs);
 
     struct Stats
     {
@@ -766,6 +766,7 @@ std::map<StorePath, StorePath> copyPaths(ref<Store> srcStore, ref<Store> dstStor
     RepairFlag repair = NoRepair,
     CheckSigsFlag checkSigs = CheckSigs,
     SubstituteFlag substitute = NoSubstitute);
+
 std::map<StorePath, StorePath> copyPaths(ref<Store> srcStore, ref<Store> dstStore,
     const StorePathSet & paths,
     RepairFlag repair = NoRepair,
diff --git a/src/libstore/worker-protocol.hh b/src/libstore/worker-protocol.hh
index e89183d40..b1888bf01 100644
--- a/src/libstore/worker-protocol.hh
+++ b/src/libstore/worker-protocol.hh
@@ -55,6 +55,7 @@ typedef enum {
     wopQueryDerivationOutputMap = 41,
     wopRegisterDrvOutput = 42,
     wopQueryRealisation = 43,
+    wopImportPaths2 = 44, // hack
 } WorkerOp;