From eeaf5ff9095d4269cb98dbbe9fe9f89d4adce157 Mon Sep 17 00:00:00 2001
From: eldritch horrors <pennae@lix.systems>
Date: Tue, 25 Feb 2025 22:13:19 +0100
Subject: [PATCH] libstore: add an async withFramedSink

Change-Id: I565606ee6e1a9cbd08bde930a1ddb07bf6ea56ef
---
 lix/libstore/remote-store-connection.hh | 13 +++++++
 lix/libstore/remote-store.cc            | 52 +++++++++++++++----------
 2 files changed, 45 insertions(+), 20 deletions(-)

diff --git a/lix/libstore/remote-store-connection.hh b/lix/libstore/remote-store-connection.hh
index 6870e878c..38a77d643 100644
--- a/lix/libstore/remote-store-connection.hh
+++ b/lix/libstore/remote-store-connection.hh
@@ -119,6 +119,19 @@ struct RemoteStore::ConnectionHandle
     void processStderr(Sink * sink = 0, Source * source = 0, bool flush = true);
 
     void withFramedSink(std::function<void(Sink & sink)> fun);
+    kj::Promise<Result<void>>
+    withFramedSinkAsync(std::function<kj::Promise<Result<void>>(Sink & sink)> fun);
+
+private:
+    struct FramedSinkHandler
+    {
+        std::exception_ptr ex;
+        std::thread stderrThread;
+
+        explicit FramedSinkHandler(ConnectionHandle & conn);
+
+        ~FramedSinkHandler() noexcept(false);
+    };
 };
 
 }
diff --git a/lix/libstore/remote-store.cc b/lix/libstore/remote-store.cc
index 6b1321d2b..0d9cbaa63 100644
--- a/lix/libstore/remote-store.cc
+++ b/lix/libstore/remote-store.cc
@@ -1020,41 +1020,53 @@ std::exception_ptr RemoteStore::Connection::processStderr(Sink * sink, Source *
     return nullptr;
 }
 
-void RemoteStore::ConnectionHandle::withFramedSink(std::function<void(Sink & sink)> fun)
+RemoteStore::ConnectionHandle::FramedSinkHandler::FramedSinkHandler(ConnectionHandle & conn)
 {
-    (*this)->to.flush();
+    conn.handle->to.flush();
 
-    std::exception_ptr ex;
-
-    /* Handle log messages / exceptions from the remote on a separate
-       thread. */
-    std::thread stderrThread([&]()
+    stderrThread = std::thread([&]()
     {
         setCurrentThreadName("remote stderr thread");
         try {
             ReceiveInterrupts receiveInterrupts;
-            processStderr(nullptr, nullptr, false);
+            conn.processStderr(nullptr, nullptr, false);
         } catch (...) {
             ex = std::current_exception();
         }
     });
+}
 
-    Finally joinStderrThread([&]()
-    {
-        if (stderrThread.joinable()) {
-            stderrThread.join();
-        }
-    });
+RemoteStore::ConnectionHandle::FramedSinkHandler::~FramedSinkHandler() noexcept(false)
+{
+    if (stderrThread.joinable()) {
+        stderrThread.join();
+    }
+    if (ex) {
+        std::rethrow_exception(ex);
+    }
+}
 
+void RemoteStore::ConnectionHandle::withFramedSink(std::function<void(Sink & sink)> fun)
+{
+    FramedSinkHandler handler{*this};
+    FramedSink sink((*this)->to, handler.ex);
+    fun(sink);
+    sink.flush();
+}
+
+kj::Promise<Result<void>> RemoteStore::ConnectionHandle::withFramedSinkAsync(
+    std::function<kj::Promise<Result<void>>(Sink & sink)> fun
+)
+try {
     {
-        FramedSink sink((*this)->to, ex);
-        fun(sink);
+        FramedSinkHandler handler{*this};
+        FramedSink sink((*this)->to, handler.ex);
+        TRY_AWAIT(fun(sink));
         sink.flush();
     }
-
-    stderrThread.join();
-    if (ex)
-        std::rethrow_exception(ex);
+    co_return result::success();
+} catch (...) {
+    co_return result::current_exception();
 }
 
 }