From 8b1d65bebe5af8960ba813e1233f2596a3ffebb7 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <edolstra@gmail.com>
Date: Tue, 14 Mar 2017 15:03:53 +0100
Subject: 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.
---
 src/libstore/binary-cache-store.cc    |  1 +
 src/libstore/download.cc              | 16 +++++++++++-----
 src/libstore/download.hh              |  3 +++
 src/libstore/s3-binary-cache-store.cc | 28 +++++++++++++++++++++++++---
 4 files changed, 40 insertions(+), 8 deletions(-)

(limited to 'src/libstore')

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
-- 
cgit v1.2.3