S3BinaryCacheStore: Support compression of narinfo and log files

You can now set the store parameter "text-compression=br" to compress
textual files in the binary cache (i.e. narinfo and logs) using
Brotli. This sets the Content-Encoding header; the extension of
compressed files is unchanged.

You can separately specify the compression of log files using
"log-compression=br". This is useful when you don't want to compress
narinfo files for backward compatibility.

5 files changed, 71 insertions, 8 deletions

diff --git a/src/libstore/binary-cache-store.cc b/src/libstore/binary-cache-store.cc
index 120345b26..804e3f6aa 100644
--- a/src/libstore/binary-cache-store.cc
+++ b/src/libstore/binary-cache-store.cc
@@ -250,6 +250,7 @@ void BinaryCacheStore::addToStore(const ValidPathInfo & info, const ref<std::str
     narInfo->url = "nar/" + printHash32(narInfo->fileHash) + ".nar"
         + (compression == "xz" ? ".xz" :
            compression == "bzip2" ? ".bz2" :
+           compression == "br" ? ".br" :
            "");
     if (repair || !fileExists(narInfo->url)) {
         stats.narWrite++;
diff --git a/src/libstore/download.cc b/src/libstore/download.cc
index d9b8fbc08..da29b2fc6 100644
--- a/src/libstore/download.cc
+++ b/src/libstore/download.cc
@@ -39,6 +39,16 @@ std::string resolveUri(const std::string & uri)
         return uri;
 }
 
+ref<std::string> decodeContent(const std::string & encoding, ref<std::string> data)
+{
+    if (encoding == "")
+        return data;
+    else if (encoding == "br")
+        return decompress(encoding, *data);
+    else
+        throw Error("unsupported Content-Encoding ‘%s’", encoding);
+}
+
 struct CurlDownloader : public Downloader
 {
     CURLM * curlm = 0;
@@ -275,12 +285,8 @@ struct CurlDownloader : public Downloader
                 result.cached = httpStatus == 304;
                 done = true;
 
-                /* Ad hoc support for brotli, since curl doesn't do
-                   this yet. */
                 try {
-                    if (encoding == "br")
-                        result.data = decompress("br", *result.data);
-
+                    result.data = decodeContent(encoding, ref<std::string>(result.data));
                     callSuccess(success, failure, const_cast<const DownloadResult &>(result));
                 } catch (...) {
                     done = true;
diff --git a/src/libstore/download.hh b/src/libstore/download.hh
index bdb5011e7..e2e16b361 100644
--- a/src/libstore/download.hh
+++ b/src/libstore/download.hh
@@ -73,4 +73,7 @@ public:
 
 bool isUri(const string & s);
 
+/* Decode data according to the Content-Encoding header. */
+ref<std::string> decodeContent(const std::string & encoding, ref<std::string> data);
+
 }
diff --git a/src/libstore/s3-binary-cache-store.cc b/src/libstore/s3-binary-cache-store.cc
index 5134dd175..1d44e6832 100644
--- a/src/libstore/s3-binary-cache-store.cc
+++ b/src/libstore/s3-binary-cache-store.cc
@@ -5,6 +5,8 @@
 #include "nar-info.hh"
 #include "nar-info-disk-cache.hh"
 #include "globals.hh"
+#include "compression.hh"
+#include "download.hh"
 
 #include <aws/core/Aws.h>
 #include <aws/core/client/ClientConfiguration.h>
@@ -104,8 +106,10 @@ S3Helper::DownloadResult S3Helper::getObject(
         auto result = checkAws(fmt("AWS error fetching ‘%s’", key),
             client->GetObject(request));
 
-        res.data = std::make_shared<std::string>(
-            dynamic_cast<std::stringstream &>(result.GetBody()).str());
+        res.data = decodeContent(
+            result.GetContentEncoding(),
+            make_ref<std::string>(
+                dynamic_cast<std::stringstream &>(result.GetBody()).str()));
 
     } catch (S3Error & e) {
         if (e.err != Aws::S3::S3Errors::NO_SUCH_KEY) throw;
@@ -137,11 +141,15 @@ struct S3BinaryCacheStoreImpl : public S3BinaryCacheStore
 
     S3Helper s3Helper;
 
+    std::string textCompression, logCompression;
+
     S3BinaryCacheStoreImpl(
         const Params & params, const std::string & bucketName)
         : S3BinaryCacheStore(params)
         , bucketName(bucketName)
         , s3Helper(get(params, "aws-region", Aws::Region::US_EAST_1))
+        , textCompression(get(params, "text-compression", "gzip"))
+        , logCompression(get(params, "log-compression", textCompression))
     {
         diskCache = getNarInfoDiskCache();
     }
@@ -220,13 +228,17 @@ struct S3BinaryCacheStoreImpl : public S3BinaryCacheStore
         return true;
     }
 
-    void upsertFile(const std::string & path, const std::string & data) override
+    void uploadFile(const std::string & path, const std::string & data,
+        const std::string & contentEncoding)
     {
         auto request =
             Aws::S3::Model::PutObjectRequest()
             .WithBucket(bucketName)
             .WithKey(path);
 
+        if (contentEncoding != "")
+            request.SetContentEncoding(contentEncoding);
+
         auto stream = std::make_shared<istringstream_nocopy>(data);
 
         request.SetBody(stream);
@@ -249,6 +261,16 @@ struct S3BinaryCacheStoreImpl : public S3BinaryCacheStore
         stats.putTimeMs += duration;
     }
 
+    void upsertFile(const std::string & path, const std::string & data) override
+    {
+        if (path.find(".narinfo") != std::string::npos)
+            uploadFile(path, *compress(textCompression, data), textCompression);
+        else if (path.find("/log") != std::string::npos)
+            uploadFile(path, *compress(logCompression, data), logCompression);
+        else
+            uploadFile(path, data, "");
+    }
+
     void getFile(const std::string & path,
         std::function<void(std::shared_ptr<std::string>)> success,
         std::function<void(std::exception_ptr exc)> failure) override
diff --git a/src/libutil/compression.cc b/src/libutil/compression.cc
index 8cb1dde66..5df97e739 100644
--- a/src/libutil/compression.cc
+++ b/src/libutil/compression.cc
@@ -91,6 +91,7 @@ static ref<std::string> decompressBzip2(const std::string & in)
 
 static ref<std::string> decompressBrotli(const std::string & in)
 {
+    // FIXME: use libbrotli
     return make_ref<std::string>(runProgram(BRO, true, {"-d"}, in));
 }
 
@@ -266,6 +267,34 @@ struct BzipSink : CompressionSink
     }
 };
 
+struct BrotliSink : CompressionSink
+{
+    Sink & nextSink;
+    std::string data;
+
+    BrotliSink(Sink & nextSink) : nextSink(nextSink)
+    {
+    }
+
+    ~BrotliSink()
+    {
+    }
+
+    // FIXME: use libbrotli
+
+    void finish() override
+    {
+        flush();
+        nextSink(runProgram(BRO, true, {}, data));
+    }
+
+    void write(const unsigned char * data, size_t len) override
+    {
+        checkInterrupt();
+        this->data.append((const char *) data, len);
+    }
+};
+
 ref<CompressionSink> makeCompressionSink(const std::string & method, Sink & nextSink)
 {
     if (method == "none")
@@ -274,6 +303,8 @@ ref<CompressionSink> makeCompressionSink(const std::string & method, Sink & next
         return make_ref<XzSink>(nextSink);
     else if (method == "bzip2")
         return make_ref<BzipSink>(nextSink);
+    else if (method == "br")
+        return make_ref<BrotliSink>(nextSink);
     else
         throw UnknownCompressionMethod(format("unknown compression method ‘%s’") % method);
 }