libutil: add async serialization helpers

Change-Id: I5c123e1ac31172d61c9e1d293e99ca1022bcec4e

lix/libstore/serve-protocol.hh

@@ -2,6 +2,7 @@
 ///@file
 
 #include "lix/libstore/common-protocol.hh"
+#include "lix/libutil/serialise-async.hh"
 
 namespace nix {
 
@@ -94,6 +95,18 @@ struct ServeProto
     {
         return ServeProto::Serialise<T>::write(conn, t);
     }
+
+    /**
+     * Create a `ServeProto::ReadConn` using the async input stream `from` and pass
+     * it to `fn`. `fn` will be run asynchronously on a fresh stack using kj fibers
+     * and can thus safely use synchronous deserializers with very little overhead.
+     */
+    static auto readAsync(auto & from, Store & store, ServeProto::Version version, auto fn)
+    {
+        return deserializeFrom(from, [&store, version, fn{std::move(fn)}](Source & wrapped) {
+            return fn(ServeProto::ReadConn{wrapped, store, version});
+        });
+    }
 };
 
 enum struct ServeProto::Command : uint64_t

lix/libstore/worker-protocol.hh

@@ -2,7 +2,9 @@
 ///@file
 
 #include "lix/libstore/common-protocol.hh"
+#include "lix/libutil/serialise-async.hh"
 #include "path-info.hh"
+#include <kj/async.h>
 
 namespace nix {
 
@@ -139,6 +141,18 @@ struct WorkerProto
     {
         return WorkerProto::Serialise<T>::write(conn, t);
     }
+
+    /**
+     * Create a `WorkerProto::ReadConn` from the async input stream `from` and pass
+     * it to `fn`. `fn` will be run asynchronously on a fresh stack using kj fibers
+     * and can thus safely use synchronous deserializers with very little overhead.
+     */
+    static auto readAsync(auto & from, Store & store, WorkerProto::Version version, auto fn)
+    {
+        return deserializeFrom(from, [&store, version, fn{std::move(fn)}](Source & wrapped) {
+            return fn(WorkerProto::ReadConn{wrapped, store, version});
+        });
+    }
 };
 
 enum struct WorkerProto::Op : uint64_t

lix/libutil/meson.build

@@ -33,6 +33,7 @@ libutil_sources = files(
   'processes.cc',
   'references.cc',
   'regex.cc',
+  'serialise-async.cc',
   'serialise.cc',
   'shlex.cc',
   'signals.cc',
@@ -118,6 +119,7 @@ libutil_headers = files(
   'result.hh',
   'rpc-fwd.hh',
   'rpc.hh',
+  'serialise-async.hh',
   'serialise.hh',
   'shlex.hh',
   'signals.hh',

lix/libutil/serialise-async.cc

@@ -0,0 +1,27 @@
+#include "serialise-async.hh"
+
+namespace nix {
+size_t detail::UnbufferedAsyncSource::read(char * data, size_t len)
+{
+    if (auto got = from.read(data, len).wait(ws).value(); got) {
+        return *got;
+    } else {
+        throw EndOfFile("async stream ended");
+    }
+}
+
+size_t detail::BufferedAsyncSource::read(char * data, size_t len)
+{
+    auto & buf = from.getBuffer();
+    if (auto avail = buf.getReadBuffer(); !avail.empty()) {
+        len = std::min(len, avail.size());
+        memcpy(data, avail.data(), len);
+        buf.consumed(len);
+        return len;
+    } else if (auto got = from.read(data, len).wait(ws).value(); got) {
+        return *got;
+    } else {
+        throw EndOfFile("async stream ended");
+    }
+}
+}

lix/libutil/serialise-async.hh

@@ -0,0 +1,79 @@
+#pragma once
+///@file Helpers for processing legacy wire protocol data on async streams
+
+#include "async-io.hh"
+#include "async.hh"
+#include "result.hh"
+#include <concepts>
+#include <kj/async.h>
+#include <type_traits>
+#include <utility>
+
+namespace nix {
+// Source wrappers for async streams. we must do this because the async deserialization overhead is
+// too large otherwise; every await or blockOn consumes far more time than the actual copy/decoding
+// done by the deserializer. this is especially important for buffered input streams since they can
+// support many small wire protocol reads on a single syscall, making the async scheduling overhead
+// even more of a loss compared to the old synchronous code. this will at least get us pretty close
+namespace detail {
+// naively adapt an async stream into a Source
+struct UnbufferedAsyncSource : Source
+{
+    kj::WaitScope & ws;
+    AsyncInputStream & from;
+
+    UnbufferedAsyncSource(kj::WaitScope & ws, AsyncInputStream & from) : ws(ws), from(from) {}
+
+    size_t read(char * data, size_t len) override;
+};
+
+// adapt a buffered async stream into a Source. unlike the unbuffered variant we will try to use the
+// read buffer as much as possible since each wait operation we do not need for IO is pure overhead.
+struct BufferedAsyncSource : Source
+{
+    kj::WaitScope & ws;
+    AsyncBufferedInputStream & from;
+
+    BufferedAsyncSource(kj::WaitScope & ws, AsyncBufferedInputStream & from) : ws(ws), from(from) {}
+
+    size_t read(char * data, size_t len) override;
+};
+
+// stacks for wrappers. the wrapper sources need wait scopes to work, and those
+// we can only get from fibers or running at the top level of an async tree. we
+// can do the latter in the daemon, but remote stores also need to deserialize.
+inline thread_local kj::FiberPool serializerFibers{65536};
+}
+
+/**
+ * Wrap the async input stream `from` in a synchronous Source and run `fn` with
+ * the wrapper as an argument, asynchronously, as a kj fiber. `fn` does not run
+ * on the main stack and instead has only 64 kiB of stack space available. `fn`
+ * should never block since only reading data from the wrapper source can yield
+ * the executor to other promises. Use async deserializers instead if possible;
+ * use this wrapper only to avoid async deserialization overhead when it hurts.
+ */
+inline auto deserializeFrom(std::derived_from<AsyncInputStream> auto & from, auto fn)
+{
+    using ResultT = decltype(fn(std::declval<Source &>()));
+    using WrapperSourceT = std::conditional_t<requires(kj::WaitScope ws) {
+        detail::BufferedAsyncSource{ws, from};
+    }, detail::BufferedAsyncSource, detail::UnbufferedAsyncSource>;
+
+    return detail::serializerFibers.startFiber(
+        [&from, fn{std::move(fn)}](kj::WaitScope & ws) -> Result<ResultT> {
+            try {
+                WrapperSourceT wrapped{ws, from};
+                if constexpr (std::is_void_v<ResultT>) {
+                    fn(wrapped);
+                    return result::success();
+                } else {
+                    return fn(wrapped);
+                }
+            } catch (...) {
+                return result::current_exception();
+            }
+        }
+    );
+}
+}