From 829af22759b8a99c3b44697365390a945f3acc04 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <eelco.dolstra@logicblox.com>
Date: Tue, 10 Jun 2014 13:30:09 +0200
Subject: [PATCH] Print a warning when loading a large path into memory

I.e. if you have a derivation with

  src = ./huge-directory;

you'll get a warning that this is not a good idea.
---
 src/libstore/remote-store.cc |  3 +++
 src/libutil/serialise.cc     | 27 +++++++++++++++++++++++++++
 src/libutil/serialise.hh     |  7 +++----
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/src/libstore/remote-store.cc b/src/libstore/remote-store.cc
index 4619206932..177b8e7afd 100644
--- a/src/libstore/remote-store.cc
+++ b/src/libstore/remote-store.cc
@@ -402,7 +402,10 @@ Path RemoteStore::addToStore(const Path & _srcPath,
     writeInt((hashAlgo == htSHA256 && recursive) ? 0 : 1, to);
     writeInt(recursive ? 1 : 0, to);
     writeString(printHashType(hashAlgo), to);
+    to.written = 0;
+    to.warn = true;
     dumpPath(srcPath, to, filter);
+    to.warn = false;
     processStderr();
     return readStorePath(from);
 }
diff --git a/src/libutil/serialise.cc b/src/libutil/serialise.cc
index 6b71f52c15..9241750750 100644
--- a/src/libutil/serialise.cc
+++ b/src/libutil/serialise.cc
@@ -54,8 +54,24 @@ FdSink::~FdSink()
 }
 
 
+size_t threshold = 256 * 1024 * 1024;
+
+static void warnLargeDump()
+{
+    printMsg(lvlError, "warning: dumping very large path (> 256 MiB); this may run out of memory");
+}
+
+
 void FdSink::write(const unsigned char * data, size_t len)
 {
+    static bool warned = false;
+    if (warn && !warned) {
+        written += len;
+        if (written > threshold) {
+            warnLargeDump();
+            warned = true;
+        }
+    }
     writeFull(fd, data, len);
 }
 
@@ -256,4 +272,15 @@ template Paths readStrings(Source & source);
 template PathSet readStrings(Source & source);
 
 
+void StringSink::operator () (const unsigned char * data, size_t len)
+{
+    static bool warned = false;
+    if (!warned && s.size() > threshold) {
+        warnLargeDump();
+        warned = true;
+    }
+    s.append((const char *) data, len);
+}
+
+
 }
diff --git a/src/libutil/serialise.hh b/src/libutil/serialise.hh
index e5a9df1d05..3d296cf6d8 100644
--- a/src/libutil/serialise.hh
+++ b/src/libutil/serialise.hh
@@ -72,6 +72,8 @@ struct BufferedSource : Source
 struct FdSink : BufferedSink
 {
     int fd;
+    bool warn;
+    size_t written;
 
     FdSink() : fd(-1) { }
     FdSink(int fd) : fd(fd) { }
@@ -95,10 +97,7 @@ struct FdSource : BufferedSource
 struct StringSink : Sink
 {
     string s;
-    void operator () (const unsigned char * data, size_t len)
-    {
-        s.append((const char *) data, len);
-    }
+    void operator () (const unsigned char * data, size_t len);
 };