From c55dcc6c13b864dc613a0a6ba51e0b897868f4b4 Mon Sep 17 00:00:00 2001
From: eldritch horrors <pennae@lix.systems>
Date: Thu, 9 May 2024 21:22:48 +0200
Subject: filetransfer: return a Source from download()

without this we will not be able to get rid of makeDecompressionSink,
which in turn will be necessary to get rid of sourceToSink (since the
libarchive archive wrapper *must* be a Source due to api limitations)

Change-Id: Iccd3d333ba2cbcab49cb5a1d3125624de16bce27
---
 src/libstore/filetransfer.cc | 120 +++++++++++++++++++++++++++----------------
 1 file changed, 77 insertions(+), 43 deletions(-)

(limited to 'src/libstore/filetransfer.cc')

diff --git a/src/libstore/filetransfer.cc b/src/libstore/filetransfer.cc
index 065f38a0c..fcb947f96 100644
--- a/src/libstore/filetransfer.cc
+++ b/src/libstore/filetransfer.cc
@@ -686,16 +686,8 @@ struct curlFileTransfer : public FileTransfer
             ->callback.get_future();
     }
 
-    void download(FileTransferRequest && request, Sink & sink) override
+    box_ptr<Source> download(FileTransferRequest && request) override
     {
-        /* Note: we can't call 'sink' via request.dataCallback, because
-           that would cause the sink to execute on the fileTransfer
-           thread. If 'sink' is a coroutine, this will fail. Also, if the
-           sink is expensive (e.g. one that does decompression and writing
-           to the Nix store), it would stall the download thread too much.
-           Therefore we use a buffer to communicate data between the
-           download thread and the calling thread. */
-
         struct State {
             bool done = false, failed = false;
             std::exception_ptr exc;
@@ -705,13 +697,6 @@ struct curlFileTransfer : public FileTransfer
 
         auto _state = std::make_shared<Sync<State>>();
 
-        /* In case of an exception, wake up the download thread. */
-        Finally finally([&]() {
-            auto state(_state->lock());
-            state->failed |= std::uncaught_exceptions() != 0;
-            state->request.notify_one();
-        });
-
         enqueueFileTransfer(
             request,
             [_state](std::exception_ptr ex) {
@@ -750,50 +735,99 @@ struct curlFileTransfer : public FileTransfer
             }
         );
 
-        std::unique_ptr<FinishSink> decompressor;
-
-        while (true) {
-            checkInterrupt();
-
+        struct InnerSource : Source
+        {
+            const std::shared_ptr<Sync<State>> _state;
             std::string chunk;
+            std::string_view buffered;
+
+            explicit InnerSource(const std::shared_ptr<Sync<State>> & state) : _state(state) {}
 
-            /* Grab data if available, otherwise wait for the download
-               thread to wake us up. */
+            ~InnerSource()
             {
+                // wake up the download thread if it's still going and have it abort
                 auto state(_state->lock());
+                state->failed |= !state->done;
+                state->request.notify_one();
+            }
 
-                if (state->data.empty()) {
-
-                    if (state->done) {
-                        if (state->exc) std::rethrow_exception(state->exc);
-                        if (decompressor) {
-                            decompressor->finish();
+            void awaitData(Sync<State>::Lock & state)
+            {
+                /* Grab data if available, otherwise wait for the download
+                   thread to wake us up. */
+                while (buffered.empty()) {
+                    if (state->data.empty()) {
+                        if (state->done) {
+                            if (state->exc) {
+                                std::rethrow_exception(state->exc);
+                            }
+                            return;
                         }
-                        return;
+
+                        state.wait(state->avail);
                     }
 
-                    state.wait(state->avail);
+                    chunk = std::move(state->data);
+                    buffered = chunk;
+                    state->request.notify_one();
+                }
+            }
 
-                    if (state->data.empty()) continue;
+            size_t read(char * data, size_t len) override
+            {
+                auto readPartial = [this](char * data, size_t len) {
+                    const auto available = std::min(len, buffered.size());
+                    memcpy(data, buffered.data(), available);
+                    buffered.remove_prefix(available);
+                    return available;
+                };
+                size_t total = readPartial(data, len);
+
+                while (total < len) {
+                    {
+                        auto state(_state->lock());
+                        awaitData(state);
+                    }
+                    const auto current = readPartial(data + total, len - total);
+                    total += current;
+                    if (total == 0 || current == 0) {
+                        break;
+                    }
+                }
+
+                if (total == 0) {
+                    throw EndOfFile("download finished");
                 }
 
-                chunk = std::move(state->data);
-                /* Reset state->data after the move, since we check data.empty() */
-                state->data = "";
+                return total;
+            }
+        };
+
+        struct DownloadSource : Source
+        {
+            InnerSource inner;
+            std::unique_ptr<Source> decompressor;
+
+            explicit DownloadSource(const std::shared_ptr<Sync<State>> & state) : inner(state) {}
+
+            size_t read(char * data, size_t len) override
+            {
+                checkInterrupt();
 
                 if (!decompressor) {
-                    decompressor = makeDecompressionSink(state->encoding, sink);
+                    auto state(inner._state->lock());
+                    inner.awaitData(state);
+                    decompressor = makeDecompressionSource(state->encoding, inner);
                 }
 
-                state->request.notify_one();
+                return decompressor->read(data, len);
             }
+        };
 
-            /* Flush the data to the sink and wake up the download thread
-               if it's blocked on a full buffer. We don't hold the state
-               lock while doing this to prevent blocking the download
-               thread if sink() takes a long time. */
-            (*decompressor)(chunk);
-        }
+        auto source = make_box_ptr<DownloadSource>(_state);
+        auto lock(_state->lock());
+        source->inner.awaitData(lock);
+        return source;
     }
 };
 
-- 
cgit v1.2.3