Merge pull request #14420 from lovesegfault/compressed-source

refactor(libutil): add `CompressedSource`

src/libstore/http-binary-cache-store.cc

@@ -141,18 +141,12 @@ void HttpBinaryCacheStore::upsertFile(
     req.method = HttpMethod::PUT;
     auto compressionMethod = getCompressionMethod(path);
 
-    std::string data;
+    std::optional<CompressedSource> compressed;
-    std::optional<StringSource> stringSource{};
 
     if (compressionMethod) {
-        StringSink sink{};
+        compressed = CompressedSource(source, *compressionMethod);
-        auto compressionSink = makeCompressionSink(*compressionMethod, sink);
-        source.drainInto(*compressionSink);
-        compressionSink->finish();
-        data = std::move(sink.s);
         req.headers.emplace_back("Content-Encoding", *compressionMethod);
-        stringSource = StringSource{data};
+        req.data = {compressed->size(), *compressed};
-        req.data = {*stringSource};
     } else {
         req.data = {sizeHint, source};
     }

src/libutil/include/nix/util/serialise.hh

@@ -272,6 +272,50 @@ struct StringSource : RestartableSource
     }
 };
 
+/**
+ * Compresses a RestartableSource using the specified compression method.
+ *
+ * @note currently this buffers the entire compressed data stream in memory. In the future it may instead compress data
+ * on demand, lazily pulling from the original `RestartableSource`. In that case, the `size()` method would go away
+ * because we would not in fact know the compressed size in advance.
+ */
+struct CompressedSource : RestartableSource
+{
+private:
+    std::string compressedData;
+    std::string compressionMethod;
+    StringSource stringSource;
+
+public:
+    /**
+     * Compress a RestartableSource using the specified compression method.
+     *
+     * @param source The source data to compress
+     * @param compressionMethod The compression method to use (e.g., "xz", "br")
+     */
+    CompressedSource(RestartableSource & source, const std::string & compressionMethod);
+
+    size_t read(char * data, size_t len) override
+    {
+        return stringSource.read(data, len);
+    }
+
+    void restart() override
+    {
+        stringSource.restart();
+    }
+
+    uint64_t size() const
+    {
+        return compressedData.size();
+    }
+
+    std::string_view getCompressionMethod() const
+    {
+        return compressionMethod;
+    }
+};
+
 /**
  * Create a restartable Source from a factory function.
  *

src/libutil/serialise.cc

@@ -1,4 +1,5 @@
 #include "nix/util/serialise.hh"
+#include "nix/util/compression.hh"
 #include "nix/util/signals.hh"
 #include "nix/util/util.hh"
 
@@ -252,6 +253,19 @@ void StringSource::skip(size_t len)
     pos += len;
 }
 
+CompressedSource::CompressedSource(RestartableSource & source, const std::string & compressionMethod)
+    : compressedData([&]() {
+        StringSink sink;
+        auto compressionSink = makeCompressionSink(compressionMethod, sink);
+        source.drainInto(*compressionSink);
+        compressionSink->finish();
+        return std::move(sink.s);
+    }())
+    , compressionMethod(compressionMethod)
+    , stringSource(compressedData)
+{
+}
+
 std::unique_ptr<FinishSink> sourceToSink(std::function<void(Source &)> fun)
 {
     struct SourceToSink : FinishSink