From fa125b9b28bea25a4eeb4d39a71a481563127cb9 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <edolstra@gmail.com>
Date: Wed, 1 Mar 2017 16:16:04 +0100
Subject: [PATCH] TeeSink: Pre-reserve string space

When receiving a very large file, this can prevent the string from
having tobe copied, which temporarily doubles memory consumption.
---
 src/libstore/export-import.cc | 11 +++++------
 src/libutil/archive.hh        | 13 +++++++++++++
 src/nix-daemon/nix-daemon.cc  |  7 +++----
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/src/libstore/export-import.cc b/src/libstore/export-import.cc
index 531f010d9..2b8ab063e 100644
--- a/src/libstore/export-import.cc
+++ b/src/libstore/export-import.cc
@@ -70,9 +70,8 @@ Paths Store::importPaths(Source & source, std::shared_ptr<FSAccessor> accessor,
         if (n != 1) throw Error("input doesn't look like something created by ‘nix-store --export’");
 
         /* Extract the NAR from the source. */
-        TeeSource tee(source);
-        ParseSink sink;
-        parseDump(sink, tee);
+        TeeSink tee(source);
+        parseDump(tee, tee.source);
 
         uint32_t magic = readInt(source);
         if (magic != exportMagic)
@@ -89,14 +88,14 @@ Paths Store::importPaths(Source & source, std::shared_ptr<FSAccessor> accessor,
         info.deriver = readString(source);
         if (info.deriver != "") assertStorePath(info.deriver);
 
-        info.narHash = hashString(htSHA256, *tee.data);
-        info.narSize = tee.data->size();
+        info.narHash = hashString(htSHA256, *tee.source.data);
+        info.narSize = tee.source.data->size();
 
         // Ignore optional legacy signature.
         if (readInt(source) == 1)
             readString(source);
 
-        addToStore(info, tee.data, false, dontCheckSigs, accessor);
+        addToStore(info, tee.source.data, false, dontCheckSigs, accessor);
 
         res.push_back(info.path);
     }
diff --git a/src/libutil/archive.hh b/src/libutil/archive.hh
index d58b91df0..c067cd2ad 100644
--- a/src/libutil/archive.hh
+++ b/src/libutil/archive.hh
@@ -70,6 +70,19 @@ struct ParseSink
     virtual void createSymlink(const Path & path, const string & target) { };
 };
 
+struct TeeSink : ParseSink
+{
+    TeeSource source;
+
+    TeeSink(Source & source) : source(source) { }
+
+    void preallocateContents(unsigned long long size) override
+    {
+        source.data->reserve(source.data->size() + size + 1024);
+    };
+
+};
+
 void parseDump(ParseSink & sink, Source & source);
 
 void restorePath(const Path & path, Source & source);
diff --git a/src/nix-daemon/nix-daemon.cc b/src/nix-daemon/nix-daemon.cc
index 174821884..ab5826b0d 100644
--- a/src/nix-daemon/nix-daemon.cc
+++ b/src/nix-daemon/nix-daemon.cc
@@ -583,12 +583,11 @@ static void performOp(ref<LocalStore> store, bool trusted, unsigned int clientVe
         if (!trusted && dontCheckSigs)
             dontCheckSigs = false;
 
-        TeeSource tee(from);
-        ParseSink sink;
-        parseDump(sink, tee);
+        TeeSink tee(from);
+        parseDump(tee, tee.source);
 
         startWork();
-        store->addToStore(info, tee.data, repair, dontCheckSigs, nullptr);
+        store->addToStore(info, tee.source.data, repair, dontCheckSigs, nullptr);
         stopWork();
         break;
     }