From 99bc9321965ced6217af2192178ec6fa94d5967e Mon Sep 17 00:00:00 2001 From: eldritch horrors Date: Mon, 3 Mar 2025 20:48:59 +0100 Subject: [PATCH] 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 | 134 +++++++++++++++++++++++++++++++ lix/libutil/async-io.hh | 140 +++++++++++++++++++++++++++++++++ lix/libutil/meson.build | 2 + tests/unit/libutil/async-io.cc | 92 ++++++++++++++++++++++ tests/unit/meson.build | 1 + 5 files changed, 369 insertions(+) create mode 100644 lix/libutil/async-io.cc create mode 100644 lix/libutil/async-io.hh create mode 100644 tests/unit/libutil/async-io.cc diff --git a/lix/libutil/async-io.cc b/lix/libutil/async-io.cc new file mode 100644 index 000000000..93eb7f9f4 --- /dev/null +++ b/lix/libutil/async-io.cc @@ -0,0 +1,134 @@ +#include "async-io.hh" + +namespace nix { +kj::Promise> AsyncInputStream::drainInto(Sink & sink) +try { + constexpr size_t BUF_SIZE = 65536; + auto buf = std::make_unique(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> AsyncInputStream::drain() +try { + StringSink s; + TRY_AWAIT(drainInto(s)); + co_return std::move(s.s); +} catch (...) { + co_return result::current_exception(); +} + +kj::Promise> AsyncSourceInputStream::read(void * buffer, size_t size) +try { + while (true) { + if (auto got = inner.read(static_cast(buffer), size); got > 0) { + return {result::success(got)}; + } + } +} catch (EndOfFile &) { + return {result::success(0)}; +} catch (...) { + return {result::current_exception()}; +} + +kj::Promise> 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> AsyncTeeInputStream::read(void * buffer, size_t size) +try { + auto got = TRY_AWAIT(inner.read(buffer, size)); + sink({static_cast(buffer), got}); + co_return got; +} catch (...) { + co_return result::current_exception(); +} + +kj::Promise> 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> 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(); + 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 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(); + 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>>> 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; +} +} diff --git a/lix/libutil/async-io.hh b/lix/libutil/async-io.hh new file mode 100644 index 000000000..dd191bb51 --- /dev/null +++ b/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 +#include +#include +#include + +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> read(void * buffer, size_t size) = 0; + + kj::Promise> drainInto(Sink & sink); + + kj::Promise> 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 owned; + +public: + AsyncSourceInputStream(Source & inner) : inner(inner) {} + AsyncSourceInputStream(box_ptr inner) : inner(*inner), owned(std::move(inner).take()) {} + + kj::Promise> 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> 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> read(void * buffer, size_t size) override; +}; + +class AsyncGeneratorInputStream : public AsyncInputStream +{ +private: + Generator g; + Bytes buf; + +public: + AsyncGeneratorInputStream(Generator && g) : g(std::move(g)) {} + + kj::Promise> 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> 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>>> result; + }; + + struct Pipe + { + // used by the source implementation + kj::Own> sendRequest; + // used by the async feeder function + kj::Promise nextRequest; + }; + + AsyncInputStream & source; + std::unique_ptr owned; + Pipe pipe; + +public: + explicit IndirectAsyncInputStreamToSource(AsyncInputStream & source); + + explicit IndirectAsyncInputStreamToSource(box_ptr 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 feed(); + + size_t read(char * data, size_t len) override; +}; + +} diff --git a/lix/libutil/meson.build b/lix/libutil/meson.build index 1e0e31d35..eb07425e1 100644 --- a/lix/libutil/meson.build +++ b/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', diff --git a/tests/unit/libutil/async-io.cc b/tests/unit/libutil/async-io.cc new file mode 100644 index 000000000..c26d67998 --- /dev/null +++ b/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 +#include +#include + +namespace nix { +TEST(IndirectAsyncInputStreamToSource, basic) +{ + struct Stream : AsyncInputStream + { + int round = 0; + kj::Promise> 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> 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(); +} +} diff --git a/tests/unit/meson.build b/tests/unit/meson.build index 4e3dfa338..cb6b13682 100644 --- a/tests/unit/meson.build +++ b/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',