libutil: add an async io header

we'll need all of these before long. we're so, so sorry.

Change-Id: I7baab54cf8a112b74d52e3a53f82836bd9f7cb83

lix/libutil/async-io.cc

@@ -0,0 +1,134 @@
+#include "async-io.hh"
+
+namespace nix {
+kj::Promise<Result<void>> AsyncInputStream::drainInto(Sink & sink)
+try {
+    constexpr size_t BUF_SIZE = 65536;
+    auto buf = std::make_unique<char[]>(BUF_SIZE);
+    while (auto r = TRY_AWAIT(read(buf.get(), BUF_SIZE))) {
+        sink(std::string_view(buf.get(), r));
+    }
+    co_return result::success();
+} catch (...) {
+    co_return result::current_exception();
+}
+
+kj::Promise<Result<std::string>> AsyncInputStream::drain()
+try {
+    StringSink s;
+    TRY_AWAIT(drainInto(s));
+    co_return std::move(s.s);
+} catch (...) {
+    co_return result::current_exception();
+}
+
+kj::Promise<Result<size_t>> AsyncSourceInputStream::read(void * buffer, size_t size)
+try {
+    while (true) {
+        if (auto got = inner.read(static_cast<char *>(buffer), size); got > 0) {
+            return {result::success(got)};
+        }
+    }
+} catch (EndOfFile &) {
+    return {result::success(0)};
+} catch (...) {
+    return {result::current_exception()};
+}
+
+kj::Promise<Result<size_t>> AsyncStringInputStream::read(void * buffer, size_t size)
+{
+    size = std::min(size, s.size());
+    if (size > 0) {
+        memcpy(buffer, s.data(), size);
+        s.remove_prefix(size);
+    }
+    return {result::success(size)};
+}
+
+kj::Promise<Result<size_t>> AsyncTeeInputStream::read(void * buffer, size_t size)
+try {
+    auto got = TRY_AWAIT(inner.read(buffer, size));
+    sink({static_cast<char *>(buffer), got});
+    co_return got;
+} catch (...) {
+    co_return result::current_exception();
+}
+
+kj::Promise<Result<size_t>> AsyncGeneratorInputStream::read(void * data, size_t len)
+try {
+    while (!buf.size()) {
+        if (auto next = g.next()) {
+            buf = *next;
+        } else {
+            return {result::success(0)};
+        }
+    }
+
+    len = std::min(len, buf.size());
+    memcpy(data, buf.data(), len);
+    buf = buf.subspan(len);
+    return {{len}};
+} catch (...) {
+    return {result::current_exception()};
+}
+
+kj::Promise<Result<size_t>> AsyncFdInputStream::read(void * buffer, size_t size)
+{
+    if (auto got = ::read(fd, buffer, size); got >= 0) {
+        return {result::success(size_t(got))};
+    } else {
+        return {result::failure(std::make_exception_ptr(SysError(errno, "read failed")))};
+    }
+}
+
+IndirectAsyncInputStreamToSource::IndirectAsyncInputStreamToSource(AsyncInputStream & source)
+    : source(source)
+    , pipe([&] {
+        auto pfp = kj::newPromiseAndCrossThreadFulfiller<Request>();
+        return Pipe{std::move(pfp.fulfiller), std::move(pfp.promise)};
+    }())
+{
+}
+
+IndirectAsyncInputStreamToSource::~IndirectAsyncInputStreamToSource() noexcept(true)
+{
+    if (pipe.sendRequest->isWaiting()) {
+        pipe.sendRequest->fulfill(Request{nullptr, 0, {}});
+    }
+}
+
+kj::Promise<void> IndirectAsyncInputStreamToSource::feed()
+{
+    while (true) {
+        auto req = co_await pipe.nextRequest;
+        if (req.data == nullptr) {
+            break;
+        }
+        try {
+            auto got = (co_await source.read(req.data, req.len)).value();
+            if (req.len != 0 && got == 0) {
+                auto eof = std::make_exception_ptr(EndOfFile("async input finished"));
+                req.result.set_exception(eof);
+                break;
+            } else {
+                auto pfp = kj::newPromiseAndCrossThreadFulfiller<Request>();
+                req.result.set_value(std::pair{got, std::move(pfp.fulfiller)});
+                pipe.nextRequest = std::move(pfp.promise);
+            }
+        } catch (...) {
+            req.result.set_exception(std::current_exception());
+            co_return;
+        }
+    }
+}
+
+size_t IndirectAsyncInputStreamToSource::read(char * data, size_t len)
+{
+    std::promise<std::pair<size_t, kj::Own<kj::CrossThreadPromiseFulfiller<Request>>>> promise;
+    auto future = promise.get_future();
+    pipe.sendRequest->fulfill(Request{data, len, std::move(promise)});
+    auto [result, next] = future.get();
+    pipe.sendRequest = std::move(next);
+    return result;
+}
+}

lix/libutil/async-io.hh

@@ -0,0 +1,140 @@
+#pragma once
+///@file
+
+#include "lix/libutil/async.hh"
+#include "lix/libutil/box_ptr.hh"
+#include "lix/libutil/result.hh"
+#include "lix/libutil/serialise.hh"
+#include <kj/async.h>
+#include <kj/common.h>
+#include <memory>
+#include <string_view>
+
+namespace nix {
+
+// not derived from kj's AsyncInputStream because read and tryRead are already
+// taken as method names, we don't need the other functions, and the bit about
+// minBytes does not work well with our current io model. some day, who knows?
+class AsyncInputStream : private kj::AsyncObject
+{
+public:
+    virtual ~AsyncInputStream() noexcept(false) {}
+
+    // expected to return 0 only on EOF or when `size = 0` was explicitly set.
+    virtual kj::Promise<Result<size_t>> read(void * buffer, size_t size) = 0;
+
+    kj::Promise<Result<void>> drainInto(Sink & sink);
+
+    kj::Promise<Result<std::string>> drain();
+};
+
+class AsyncSourceInputStream : public AsyncInputStream
+{
+    Source & inner;
+    // inner must reference owned if owned is set. we'll keep a unique_ptr
+    // field around in all instances to avoid duplicating the entire class
+    // into a reference variant and an owning variant (holding a box_ptr).
+    std::unique_ptr<Source> owned;
+
+public:
+    AsyncSourceInputStream(Source & inner) : inner(inner) {}
+    AsyncSourceInputStream(box_ptr<Source> inner) : inner(*inner), owned(std::move(inner).take()) {}
+
+    kj::Promise<Result<size_t>> read(void * buffer, size_t size) override;
+};
+
+class AsyncStringInputStream : public AsyncInputStream
+{
+    std::string_view s;
+
+public:
+    explicit AsyncStringInputStream(std::string_view s) : s(s) {}
+
+    kj::Promise<Result<size_t>> read(void * buffer, size_t size) override;
+};
+
+// this writes to sources instead of async streams because none of the sinks
+// we need to date are actually async, not because that wouldn't be possible
+class AsyncTeeInputStream : public AsyncInputStream
+{
+    AsyncInputStream & inner;
+    Sink & sink;
+
+public:
+    AsyncTeeInputStream(AsyncInputStream & inner, Sink & sink) : inner(inner), sink(sink) {}
+
+    kj::Promise<Result<size_t>> read(void * buffer, size_t size) override;
+};
+
+class AsyncGeneratorInputStream : public AsyncInputStream
+{
+private:
+    Generator<Bytes> g;
+    Bytes buf;
+
+public:
+    AsyncGeneratorInputStream(Generator<Bytes> && g) : g(std::move(g)) {}
+
+    kj::Promise<Result<size_t>> read(void * data, size_t len) override;
+};
+
+class AsyncFdInputStream : public AsyncInputStream
+{
+    int fd;
+    AutoCloseFD ownedFd; // only for closing automatically, must equal fd if set
+
+public:
+    struct shared_fd
+    {};
+
+    explicit AsyncFdInputStream(AutoCloseFD fd) : fd(fd.get()), ownedFd(std::move(fd)) {}
+    AsyncFdInputStream(shared_fd, int fd) : fd(fd) {}
+
+    kj::Promise<Result<size_t>> read(void * buffer, size_t size) override;
+};
+
+/**
+ * Wraps a stream in a source. The returned source must not be used on the
+ * event loop that created it, otherwise read requests cannot be serviced.
+ */
+class IndirectAsyncInputStreamToSource : public Source
+{
+    struct Request
+    {
+        char * data;
+        size_t len;
+        std::promise<std::pair<size_t, kj::Own<kj::CrossThreadPromiseFulfiller<Request>>>> result;
+    };
+
+    struct Pipe
+    {
+        // used by the source implementation
+        kj::Own<kj::CrossThreadPromiseFulfiller<Request>> sendRequest;
+        // used by the async feeder function
+        kj::Promise<Request> nextRequest;
+    };
+
+    AsyncInputStream & source;
+    std::unique_ptr<AsyncInputStream> owned;
+    Pipe pipe;
+
+public:
+    explicit IndirectAsyncInputStreamToSource(AsyncInputStream & source);
+
+    explicit IndirectAsyncInputStreamToSource(box_ptr<AsyncInputStream> owned)
+        : IndirectAsyncInputStreamToSource(*owned)
+    {
+        this->owned = std::move(owned).take();
+    }
+
+    ~IndirectAsyncInputStreamToSource() noexcept(true);
+
+    KJ_DISALLOW_COPY_AND_MOVE(IndirectAsyncInputStreamToSource);
+
+    /** Feed the source. Must be awaited fully to drain the input stream. */
+    kj::Promise<void> feed();
+
+    size_t read(char * data, size_t len) override;
+};
+
+}

lix/libutil/meson.build

@@ -1,6 +1,7 @@
 libutil_sources = files(
   'archive.cc',
   'args.cc',
+  'async-io.cc',
   'canon-path.cc',
   'cgroup.cc',
   'compression.cc',
@@ -55,6 +56,7 @@ libutil_headers = files(
   'args/root.hh',
   'args.hh',
   'async-collect.hh',
+  'async-io.hh',
   'async.hh',
   'async-semaphore.hh',
   'backed-string-view.hh',

tests/unit/libutil/async-io.cc

@@ -0,0 +1,92 @@
+#include "lix/libutil/async-io.hh"
+#include "lix/libutil/async.hh"
+#include "lix/libutil/file-descriptor.hh"
+#include "lix/libutil/result.hh"
+
+#include <exception>
+#include <gtest/gtest.h>
+#include <stdexcept>
+
+namespace nix {
+TEST(IndirectAsyncInputStreamToSource, basic)
+{
+    struct Stream : AsyncInputStream
+    {
+        int round = 0;
+        kj::Promise<Result<size_t>> read(void * buffer, size_t size) override
+        {
+            round++;
+            if (round <= 10) {
+                memset(buffer, size, 1);
+                return {{1}};
+            } else if (round <= 13) {
+                memset(buffer, size, size);
+                return {{size}};
+            } else {
+                return {result::success(0)};
+            }
+        }
+    };
+
+    AsyncIoRoot aio;
+    Stream s;
+
+    IndirectAsyncInputStreamToSource is(s);
+
+    auto user = std::async(std::launch::async, [&]() {
+        char buf[1026];
+
+        // single read
+        ASSERT_EQ(is.read(buf, 1), 1);
+        ASSERT_EQ(buf[0], 1);
+
+        // read spanning blocks doesn't coalesce
+        ASSERT_EQ(is.read(buf, 5), 1);
+        ASSERT_EQ(buf[0], 5);
+
+        // coalescing from Source works
+        ASSERT_NO_THROW(is(buf, 8));
+        ASSERT_EQ(buf[0], 8);
+
+        // next reads fill all sizes
+        ASSERT_EQ(is.read(buf, 513), 513);
+        ASSERT_EQ(buf[0], 1);
+        ASSERT_EQ(is.read(buf, 1025), 1025);
+        ASSERT_EQ(buf[0], 1);
+
+        // zero-size reads don't EOF
+        ASSERT_EQ(is.read(buf, 0), 0);
+
+        // EOF propagates
+        ASSERT_THROW(is.read(buf, 1025), EndOfFile);
+    });
+
+    is.feed().wait(aio.kj.waitScope);
+    user.get();
+}
+
+TEST(IndirectAsyncInputStreamToSource, errorPropagation)
+{
+    struct Stream : AsyncInputStream
+    {
+        int round = 0;
+        kj::Promise<Result<size_t>> read(void * buffer, size_t size) override
+        {
+            return {result::failure(std::make_exception_ptr(std::invalid_argument("foo")))};
+        }
+    };
+
+    AsyncIoRoot aio;
+    Stream s;
+
+    IndirectAsyncInputStreamToSource is(s);
+
+    auto user = std::async(std::launch::async, [&]() {
+        char buf[1];
+        ASSERT_THROW(is.read(buf, 1), std::invalid_argument);
+    });
+
+    is.feed().wait(aio.kj.waitScope);
+    user.get();
+}
+}

tests/unit/meson.build

@@ -44,6 +44,7 @@ liblixutil_test_support = declare_dependency(
 libutil_tests_sources = files(
   'libutil/archive.cc',
   'libutil/async-collect.cc',
+  'libutil/async-io.cc',
   'libutil/async-semaphore.cc',
   'libutil/canon-path.cc',
   'libutil/checked-arithmetic.cc',