libstore: rewrite narFromPath as generator

Change-Id: Ifa783c2c65c06ddd1d0212016d5bfd07666ea91c

18 files changed, 47 insertions, 36 deletions

diff --git a/src/libstore/binary-cache-store.cc b/src/libstore/binary-cache-store.cc
index 27ea621db..e507391fe 100644
--- a/src/libstore/binary-cache-store.cc
+++ b/src/libstore/binary-cache-store.cc
@@ -329,25 +329,32 @@ std::optional<StorePath> BinaryCacheStore::queryPathFromHashPart(const std::stri
     }
 }
 
-void BinaryCacheStore::narFromPath(const StorePath & storePath, Sink & sink)
+WireFormatGenerator BinaryCacheStore::narFromPath(const StorePath & storePath)
 {
     auto info = queryPathInfo(storePath).cast<const NarInfo>();
 
-    LengthSink narSize;
-    TeeSink tee { sink, narSize };
-
-
     try {
         auto file = getFile(info->url);
-        auto decompressor = makeDecompressionSource(info->compression, *file);
-        decompressor->drainInto(tee);
+        return [](auto info, auto file, auto & stats) -> WireFormatGenerator {
+            std::unique_ptr<char[]> buf(new char[65536]);
+            size_t total = 0;
+            auto decompressor = makeDecompressionSource(info->compression, *file);
+            try {
+                while (true) {
+                    const auto len = decompressor->read(buf.get(), sizeof(buf));
+                    co_yield std::span{buf.get(), len};
+                    total += len;
+                }
+            } catch (EndOfFile &) {
+            }
+
+            stats.narRead++;
+            //stats.narReadCompressedBytes += nar->size(); // FIXME
+            stats.narReadBytes += total;
+        }(std::move(info), std::move(file), stats);
     } catch (NoSuchBinaryCacheFile & e) {
         throw SubstituteGone(std::move(e.info()));
     }
-
-    stats.narRead++;
-    //stats.narReadCompressedBytes += nar->size(); // FIXME
-    stats.narReadBytes += narSize.length;
 }
 
 std::shared_ptr<const ValidPathInfo> BinaryCacheStore::queryPathInfoUncached(const StorePath & storePath)
diff --git a/src/libstore/binary-cache-store.hh b/src/libstore/binary-cache-store.hh
index 6c1f1348c..2ad26cc65 100644
--- a/src/libstore/binary-cache-store.hh
+++ b/src/libstore/binary-cache-store.hh
@@ -136,7 +136,7 @@ public:
 
     std::shared_ptr<const Realisation> queryRealisationUncached(const DrvOutput &) override;
 
-    void narFromPath(const StorePath & path, Sink & sink) override;
+    WireFormatGenerator narFromPath(const StorePath & path) override;
 
     ref<FSAccessor> getFSAccessor() override;
 
diff --git a/src/libstore/build/local-derivation-goal.cc b/src/libstore/build/local-derivation-goal.cc
index a071883bb..5ca8eac95 100644
--- a/src/libstore/build/local-derivation-goal.cc
+++ b/src/libstore/build/local-derivation-goal.cc
@@ -1085,11 +1085,11 @@ struct RestrictedStore : public virtual RestrictedStoreConfig, public virtual In
         return path;
     }
 
-    void narFromPath(const StorePath & path, Sink & sink) override
+    WireFormatGenerator narFromPath(const StorePath & path) override
     {
         if (!goal.isAllowed(path))
             throw InvalidPath("cannot dump unknown path '%s' in recursive Nix", printStorePath(path));
-        LocalFSStore::narFromPath(path, sink);
+        return LocalFSStore::narFromPath(path);
     }
 
     void ensurePath(const StorePath & path) override
diff --git a/src/libstore/dummy-store.cc b/src/libstore/dummy-store.cc
index e97425c80..e16f87e4b 100644
--- a/src/libstore/dummy-store.cc
+++ b/src/libstore/dummy-store.cc
@@ -63,7 +63,7 @@ struct DummyStore : public virtual DummyStoreConfig, public virtual Store
         RepairFlag repair) override
     { unsupported("addTextToStore"); }
 
-    void narFromPath(const StorePath & path, Sink & sink) override
+    WireFormatGenerator narFromPath(const StorePath & path) override
     { unsupported("narFromPath"); }
 
     std::shared_ptr<const Realisation> queryRealisationUncached(const DrvOutput &) override
diff --git a/src/libstore/export-import.cc b/src/libstore/export-import.cc
index a05a9f23b..fa1020537 100644
--- a/src/libstore/export-import.cc
+++ b/src/libstore/export-import.cc
@@ -33,7 +33,7 @@ void Store::exportPath(const StorePath & path, Sink & sink)
     HashSink hashSink(htSHA256);
     TeeSink teeSink(sink, hashSink);
 
-    narFromPath(path, teeSink);
+    teeSink << narFromPath(path);
 
     /* Refuse to export paths that have changed.  This prevents
        filesystem corruption from spreading to other machines.
diff --git a/src/libstore/legacy-ssh-store.cc b/src/libstore/legacy-ssh-store.cc
index b4f3854d5..57d333f12 100644
--- a/src/libstore/legacy-ssh-store.cc
+++ b/src/libstore/legacy-ssh-store.cc
@@ -227,13 +227,15 @@ struct LegacySSHStore : public virtual LegacySSHStoreConfig, public virtual Stor
             throw Error("failed to add path '%s' to remote host '%s'", printStorePath(info.path), host);
     }
 
-    void narFromPath(const StorePath & path, Sink & sink) override
+    WireFormatGenerator narFromPath(const StorePath & path) override
     {
         auto conn(connections->get());
 
         conn->to << ServeProto::Command::DumpStorePath << printStorePath(path);
         conn->to.flush();
-        sink << copyNAR(conn->from);
+        return [] (auto conn) -> WireFormatGenerator {
+            co_yield copyNAR(conn->from);
+        }(std::move(conn));
     }
 
     std::optional<StorePath> queryPathFromHashPart(const std::string & hashPart) override
diff --git a/src/libstore/local-fs-store.cc b/src/libstore/local-fs-store.cc
index a16e0ba73..b5b1959be 100644
--- a/src/libstore/local-fs-store.cc
+++ b/src/libstore/local-fs-store.cc
@@ -78,11 +78,11 @@ ref<FSAccessor> LocalFSStore::getFSAccessor()
             std::dynamic_pointer_cast<LocalFSStore>(shared_from_this())));
 }
 
-void LocalFSStore::narFromPath(const StorePath & path, Sink & sink)
+WireFormatGenerator LocalFSStore::narFromPath(const StorePath & path)
 {
     if (!isValidPath(path))
         throw Error("path '%s' does not exist in store", printStorePath(path));
-    sink << dumpPath(getRealStoreDir() + std::string(printStorePath(path), storeDir.size()));
+    return dumpPath(getRealStoreDir() + std::string(printStorePath(path), storeDir.size()));
 }
 
 const std::string LocalFSStore::drvsLogDir = "drvs";
diff --git a/src/libstore/local-fs-store.hh b/src/libstore/local-fs-store.hh
index e03023916..56de47424 100644
--- a/src/libstore/local-fs-store.hh
+++ b/src/libstore/local-fs-store.hh
@@ -42,7 +42,7 @@ public:
 
     LocalFSStore(const Params & params);
 
-    void narFromPath(const StorePath & path, Sink & sink) override;
+    WireFormatGenerator narFromPath(const StorePath & path) override;
     ref<FSAccessor> getFSAccessor() override;
 
     /**
diff --git a/src/libstore/make-content-addressed.cc b/src/libstore/make-content-addressed.cc
index 1aa139b38..4e6d9dfe5 100644
--- a/src/libstore/make-content-addressed.cc
+++ b/src/libstore/make-content-addressed.cc
@@ -23,7 +23,7 @@ std::map<StorePath, StorePath> makeContentAddressed(
         std::string oldHashPart(path.hashPart());
 
         StringSink sink;
-        srcStore.narFromPath(path, sink);
+        sink << srcStore.narFromPath(path);
 
         StringMap rewrites;
 
diff --git a/src/libstore/remote-fs-accessor.cc b/src/libstore/remote-fs-accessor.cc
index 8da8b9774..0689ce74d 100644
--- a/src/libstore/remote-fs-accessor.cc
+++ b/src/libstore/remote-fs-accessor.cc
@@ -97,7 +97,7 @@ std::pair<ref<FSAccessor>, Path> RemoteFSAccessor::fetch(const Path & path_, boo
     }
 
     StringSink sink;
-    store->narFromPath(storePath, sink);
+    sink << store->narFromPath(storePath);
     return {addToCache(storePath.hashPart(), std::move(sink.s)), restPath};
 }
 
diff --git a/src/libstore/remote-store.cc b/src/libstore/remote-store.cc
index 9cebdbe82..9ff01351c 100644
--- a/src/libstore/remote-store.cc
+++ b/src/libstore/remote-store.cc
@@ -848,12 +848,14 @@ RemoteStore::Connection::~Connection()
     }
 }
 
-void RemoteStore::narFromPath(const StorePath & path, Sink & sink)
+WireFormatGenerator RemoteStore::narFromPath(const StorePath & path)
 {
     auto conn(connections->get());
     conn->to << WorkerProto::Op::NarFromPath << printStorePath(path);
     conn->processStderr();
-    sink << copyNAR(conn->from);
+    return [](auto conn) -> WireFormatGenerator {
+        co_yield copyNAR(conn->from);
+    }(std::move(conn));
 }
 
 ref<FSAccessor> RemoteStore::getFSAccessor()
diff --git a/src/libstore/remote-store.hh b/src/libstore/remote-store.hh
index b1b7f93e9..264869df6 100644
--- a/src/libstore/remote-store.hh
+++ b/src/libstore/remote-store.hh
@@ -183,7 +183,7 @@ protected:
 
     virtual ref<FSAccessor> getFSAccessor() override;
 
-    virtual void narFromPath(const StorePath & path, Sink & sink) override;
+    virtual WireFormatGenerator narFromPath(const StorePath & path) override;
 
 private:
 
diff --git a/src/libstore/store-api.cc b/src/libstore/store-api.cc
index 683da1129..28a414555 100644
--- a/src/libstore/store-api.cc
+++ b/src/libstore/store-api.cc
@@ -1065,7 +1065,7 @@ void copyStorePath(
             act.progress(total, info->narSize);
         });
         TeeSink tee { sink, progressSink };
-        srcStore.narFromPath(storePath, tee);
+        tee << srcStore.narFromPath(storePath);
     });
 
     dstStore.addToStore(*info, *source, repair, checkSigs);
@@ -1202,7 +1202,7 @@ std::map<StorePath, StorePath> copyPaths(
                 });
                 TeeSink tee{sink, progressSink};
 
-                srcStore.narFromPath(missingPath, tee);
+                tee << srcStore.narFromPath(missingPath);
             });
         pathsToCopy.push_back(std::pair{infoForDst, std::move(source)});
     }
diff --git a/src/libstore/store-api.hh b/src/libstore/store-api.hh
index 186437f43..91a8ef2ca 100644
--- a/src/libstore/store-api.hh
+++ b/src/libstore/store-api.hh
@@ -577,9 +577,9 @@ public:
     { return registerDrvOutput(output); }
 
     /**
-     * Write a NAR dump of a store path.
+     * Generate a NAR dump of a store path.
      */
-    virtual void narFromPath(const StorePath & path, Sink & sink) = 0;
+    virtual WireFormatGenerator narFromPath(const StorePath & path) = 0;
 
     /**
      * For each path, if it's a derivation, build it.  Building a
diff --git a/src/libstore/uds-remote-store.hh b/src/libstore/uds-remote-store.hh
index cdb28a001..ff7e9ae3f 100644
--- a/src/libstore/uds-remote-store.hh
+++ b/src/libstore/uds-remote-store.hh
@@ -38,8 +38,8 @@ public:
     ref<FSAccessor> getFSAccessor() override
     { return LocalFSStore::getFSAccessor(); }
 
-    void narFromPath(const StorePath & path, Sink & sink) override
-    { LocalFSStore::narFromPath(path, sink); }
+    WireFormatGenerator narFromPath(const StorePath & path) override
+    { return LocalFSStore::narFromPath(path); }
 
     /**
      * Implementation of `IndirectRootStore::addIndirectRoot()` which
diff --git a/src/nix-store/nix-store.cc b/src/nix-store/nix-store.cc
index 240888a16..f11e2fe4d 100644
--- a/src/nix-store/nix-store.cc
+++ b/src/nix-store/nix-store.cc
@@ -763,7 +763,7 @@ static void opVerifyPath(Strings opFlags, Strings opArgs)
         printMsg(lvlTalkative, "checking path '%s'...", store->printStorePath(path));
         auto info = store->queryPathInfo(path);
         HashSink sink(info->narHash.type);
-        store->narFromPath(path, sink);
+        sink << store->narFromPath(path);
         auto current = sink.finish();
         if (current.first != info->narHash) {
             printError("path '%s' was modified! expected hash '%s', got '%s'",
@@ -900,7 +900,7 @@ static void opServe(Strings opFlags, Strings opArgs)
             }
 
             case ServeProto::Command::DumpStorePath:
-                store->narFromPath(store->parseStorePath(readString(in)), out);
+                out << store->narFromPath(store->parseStorePath(readString(in)));
                 break;
 
             case ServeProto::Command::ImportPaths: {
diff --git a/src/nix/dump-path.cc b/src/nix/dump-path.cc
index b07c76136..661438a3d 100644
--- a/src/nix/dump-path.cc
+++ b/src/nix/dump-path.cc
@@ -22,7 +22,7 @@ struct CmdDumpPath : StorePathCommand
     {
         logger->pause();
         FdSink sink(STDOUT_FILENO);
-        store->narFromPath(storePath, sink);
+        sink << store->narFromPath(storePath);
         sink.flush();
     }
 };
diff --git a/src/nix/verify.cc b/src/nix/verify.cc
index 9be4c25dc..eb68e67bc 100644
--- a/src/nix/verify.cc
+++ b/src/nix/verify.cc
@@ -100,7 +100,7 @@ struct CmdVerify : StorePathsCommand
 
                     auto hashSink = HashSink(info->narHash.type);
 
-                    store->narFromPath(info->path, hashSink);
+                    hashSink << store->narFromPath(info->path);
 
                     auto hash = hashSink.finish();