libstore: add framed data support to sendCommand

the subframing layer is ... a bit of challenge. since the old code is synchronous but wants to handle errors asynchronously anyway it is on the subframing layer to *spawn a thread* that polls for errors on the wire, while non-framed commands handle errors synchronously once they have sent all their data. this encapsulation of the wires is far from perfect (let alone legible), but hopefully it will be only temporary. Change-Id: I26d8020549b767794cae121313360c488504995f
2025-06-11 22:29:30 +02:00
parent 1a2247560d
commit 8b3fdbc847
2 changed files with 72 additions and 51 deletions
@@ -7,7 +7,9 @@
 #include "lix/libutil/pool.hh"
 #include "lix/libutil/result.hh"
 #include "lix/libutil/serialise.hh"
+#include <tuple>
 #include <type_traits>
+#include <utility>

 namespace nix {

@@ -140,8 +142,29 @@ struct RemoteStore::ConnectionHandle
    template<typename R = void, typename... Args>
    kj::Promise<Result<R>> sendCommand(Args &&... args)
    try {
-        ((handle->to << std::forward<Args>(args)), ...);
-        processStderr();
+        constexpr auto LastArgIdx = sizeof...(Args) - 1;
+        using AllArgsT = std::tuple<Args &&...>;
+        using LastArgT = std::tuple_element_t<LastArgIdx, AllArgsT>;
+
+        // if the last argument can be serialized normally we will serialize *all*
+        // arguments at once and hand off to the remote. if the last argument does
+        // not have a serializer we assume it's a callback for a subframe protocol
+        // and serialize all *preceding* arguments normally before handing over to
+        // the subframing layer (which is then responsible for any error handling)
+        if constexpr (requires { handle->to << std::declval<LastArgT>(); }) {
+            ((handle->to << std::forward<Args>(args)), ...);
+            processStderr();
+        } else {
+            using ImmediateArgsIdxs = std::make_index_sequence<sizeof...(Args) - 1>;
+            AllArgsT allArgs(std::forward<Args>(args)...);
+
+            [&]<size_t... Ids>(std::integer_sequence<size_t, Ids...>) {
+                ((handle->to << std::get<Ids>(std::forward<AllArgsT>(allArgs))), ...);
+            }(ImmediateArgsIdxs{});
+
+            LIX_TRY_AWAIT(withFramedSinkAsync(std::get<LastArgIdx>(allArgs)));
+        }
+
        if constexpr (std::is_void_v<R>) {
            co_return result::success();
        } else {
@@ -354,24 +354,18 @@ kj::Promise<Result<ref<const ValidPathInfo>>> RemoteStore::addCAToStore(
 try {
    auto conn(TRY_AWAIT(getConnection()));

-    conn->to
-        << WorkerProto::Op::AddToStore
-        << name
-        << caMethod.render(hashType);
-    conn->to << WorkerProto::write(*conn, references);
-    conn->to << repair;
-
    // The dump source may invoke the store, so we need to make some room.
    connections->incCapacity();
-    {
-        Finally cleanup([&]() { connections->decCapacity(); });
-        TRY_AWAIT(conn.withFramedSinkAsync([&](Sink & sink) {
-            return dump.drainInto(sink);
-        }));
-    }
+    Finally cleanup([&]() { connections->decCapacity(); });

-    co_return make_ref<ValidPathInfo>(
-        WorkerProto::Serialise<ValidPathInfo>::read(*conn));
+    co_return make_ref<ValidPathInfo>(TRY_AWAIT(conn.sendCommand<ValidPathInfo>(
+        WorkerProto::Op::AddToStore,
+        name,
+        caMethod.render(hashType),
+        WorkerProto::write(*conn, references),
+        repair,
+        [&](Sink & sink) { return dump.drainInto(sink); }
+    )));
 } catch (...) {
    co_return result::current_exception();
 }
@@ -401,19 +395,22 @@ kj::Promise<Result<void>> RemoteStore::addToStore(
 try {
    auto conn(TRY_AWAIT(getConnection()));

-    conn->to << WorkerProto::Op::AddToStoreNar
-             << printStorePath(info.path)
-             << (info.deriver ? printStorePath(*info.deriver) : "")
-             << info.narHash.to_string(Base::Base16, false);
-    conn->to << WorkerProto::write(*conn, info.references);
-    conn->to << info.registrationTime << info.narSize
-             << info.ultimate << info.sigs << renderContentAddress(info.ca)
-             << repair << !checkSigs;
-
    auto copier = copyNAR(source);
-    TRY_AWAIT(conn.withFramedSinkAsync([&](Sink & sink) {
-        return copier->drainInto(sink);
-    }));
+    TRY_AWAIT(conn.sendCommand(
+        WorkerProto::Op::AddToStoreNar,
+        printStorePath(info.path),
+        (info.deriver ? printStorePath(*info.deriver) : ""),
+        info.narHash.to_string(Base::Base16, false),
+        WorkerProto::write(*conn, info.references),
+        info.registrationTime,
+        info.narSize,
+        info.ultimate,
+        info.sigs,
+        renderContentAddress(info.ca),
+        repair,
+        !checkSigs,
+        [&](Sink & sink) { return copier->drainInto(sink); }
+    ));
    co_return result::success();
 } catch (...) {
    co_return result::current_exception();
@@ -429,25 +426,26 @@ try {
    auto remoteVersion = TRY_AWAIT(getProtocol());

    auto conn(TRY_AWAIT(getConnection()));
-    conn->to
-        << WorkerProto::Op::AddMultipleToStore
-        << repair
-        << !checkSigs;
-    // NOLINTNEXTLINE(cppcoreguidelines-avoid-capturing-lambda-coroutines)
-    TRY_AWAIT(conn.withFramedSinkAsync([&](Sink & sink) -> kj::Promise<Result<void>> {
-        try {
-            sink << pathsToCopy.size();
-            for (auto & [pathInfo, pathSource] : pathsToCopy) {
-                sink << WorkerProto::Serialise<ValidPathInfo>::write(
-                    WorkerProto::WriteConn {*this, remoteVersion},
-                    pathInfo);
-                TRY_AWAIT(TRY_AWAIT(pathSource())->drainInto(sink));
+    TRY_AWAIT(conn.sendCommand(
+        WorkerProto::Op::AddMultipleToStore,
+        repair,
+        !checkSigs,
+        // NOLINTNEXTLINE(cppcoreguidelines-avoid-capturing-lambda-coroutines)
+        [&](Sink & sink) -> kj::Promise<Result<void>> {
+            try {
+                sink << pathsToCopy.size();
+                for (auto & [pathInfo, pathSource] : pathsToCopy) {
+                    sink << WorkerProto::Serialise<ValidPathInfo>::write(
+                        WorkerProto::WriteConn{*this, remoteVersion}, pathInfo
+                    );
+                    TRY_AWAIT(TRY_AWAIT(pathSource())->drainInto(sink));
+                }
+                co_return result::success();
+            } catch (...) {
+                co_return result::current_exception();
            }
-            co_return result::success();
-        } catch (...) {
-            co_return result::current_exception();
        }
-    }));
+    ));
    co_return result::success();
 } catch (...) {
    co_return result::current_exception();
@@ -654,12 +652,12 @@ try {
 kj::Promise<Result<void>> RemoteStore::addBuildLog(const StorePath & drvPath, std::string_view log)
 try {
    auto conn(TRY_AWAIT(getConnection()));
-    conn->to << WorkerProto::Op::AddBuildLog << drvPath.to_string();
    AsyncStringInputStream source(log);
-    TRY_AWAIT(conn.withFramedSinkAsync([&](Sink & sink) {
-        return source.drainInto(sink);
-    }));
-    readInt(conn->from);
+    TRY_AWAIT(conn.sendCommand<unsigned>(
+        WorkerProto::Op::AddBuildLog,
+        drvPath.to_string(),
+        [&](Sink & sink) { return source.drainInto(sink); }
+    ));
    co_return result::success();
 } catch (...) {
    co_return result::current_exception();