From 47f7e5585bbb9a6acfaddda7327ba4eeb5b662c2 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <edolstra@gmail.com>
Date: Wed, 21 Mar 2018 22:56:02 +0100
Subject: [PATCH] Make 'nix copy --from ssh://...' run in constant memory

For instance, this reduced the memory consumption of

  $ nix copy --from ssh://localhost --to ~/my-nix /nix/store/1n7x0yv8vq6zi90hfmian84vdhd04bgp-blender-2.79a

from 632 MiB to 16 MiB.
---
 src/libstore/legacy-ssh-store.cc |  7 +------
 src/libstore/ssh-store.cc        | 19 +------------------
 src/libutil/archive.cc           | 17 +++++++++++++++++
 src/libutil/archive.hh           |  3 +++
 4 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/src/libstore/legacy-ssh-store.cc b/src/libstore/legacy-ssh-store.cc
index 2a2e3d914..b52de5434 100644
--- a/src/libstore/legacy-ssh-store.cc
+++ b/src/libstore/legacy-ssh-store.cc
@@ -151,12 +151,7 @@ struct LegacySSHStore : public Store
 
         conn->to << cmdDumpStorePath << path;
         conn->to.flush();
-
-        /* FIXME: inefficient. */
-        ParseSink parseSink; /* null sink; just parse the NAR */
-        TeeSource savedNAR(conn->from);
-        parseDump(parseSink, savedNAR);
-        sink(*savedNAR.data);
+        copyNAR(conn->from, sink);
     }
 
     PathSet queryAllValidPaths() override { unsupported(); }
diff --git a/src/libstore/ssh-store.cc b/src/libstore/ssh-store.cc
index 398408ea8..39205ae2c 100644
--- a/src/libstore/ssh-store.cc
+++ b/src/libstore/ssh-store.cc
@@ -63,29 +63,12 @@ private:
     };
 };
 
-
-class ForwardSource : public Source
-{
-    Source & readSource;
-    Sink & writeSink;
-public:
-    ForwardSource(Source & readSource, Sink & writeSink) : readSource(readSource), writeSink(writeSink) {}
-    size_t read(unsigned char * data, size_t len) override
-    {
-        auto n = readSource.read(data, len);
-        writeSink(data, n);
-        return n;
-    }
-};
-
 void SSHStore::narFromPath(const Path & path, Sink & sink)
 {
     auto conn(connections->get());
     conn->to << wopNarFromPath << path;
     conn->processStderr();
-    ParseSink ps;
-    auto fwd = ForwardSource(conn->from, sink);
-    parseDump(ps, fwd);
+    copyNAR(conn->from, sink);
 }
 
 ref<FSAccessor> SSHStore::getFSAccessor()
diff --git a/src/libutil/archive.cc b/src/libutil/archive.cc
index f71229d8f..a1d9d3233 100644
--- a/src/libutil/archive.cc
+++ b/src/libutil/archive.cc
@@ -350,4 +350,21 @@ void restorePath(const Path & path, Source & source)
 }
 
 
+void copyNAR(Source & source, Sink & sink)
+{
+    // FIXME: if 'source' is the output of dumpPath() followed by EOF,
+    // we should just forward all data directly without parsing.
+
+    ParseSink parseSink; /* null sink; just parse the NAR */
+
+    LambdaSource wrapper([&](unsigned char * data, size_t len) {
+        auto n = source.read(data, len);
+        sink(data, n);
+        return n;
+    });
+
+    parseDump(parseSink, wrapper);
+}
+
+
 }
diff --git a/src/libutil/archive.hh b/src/libutil/archive.hh
index 8a15e849c..7a0e688e4 100644
--- a/src/libutil/archive.hh
+++ b/src/libutil/archive.hh
@@ -74,6 +74,9 @@ void parseDump(ParseSink & sink, Source & source);
 
 void restorePath(const Path & path, Source & source);
 
+/* Read a NAR from 'source' and write it to 'sink'. */
+void copyNAR(Source & source, Sink & sink);
+
 
 // FIXME: global variables are bad m'kay.
 extern bool useCaseHack;