diff --git a/lix/libutil/async-io.hh b/lix/libutil/async-io.hh index 92a28dbc6..66400271c 100644 --- a/lix/libutil/async-io.hh +++ b/lix/libutil/async-io.hh @@ -9,10 +9,12 @@ #include "lix/libutil/ref.hh" #include "lix/libutil/result.hh" #include "lix/libutil/serialise.hh" +#include #include #include #include #include +#include #include #include @@ -273,4 +275,87 @@ struct AsyncPipe * multiple readers are allowed, but are serviced in some indeterminate order. */ AsyncPipe newZeroCopyPipe(); + +/** + * Adapts an `AsyncInputStream`-consuming promise created by a callback into a + * pair of `AsyncOutputStream` and a callback that awaits the wrapped promise. + * Errors flow to both the writer *and* the promise created by the callback we + * return here. Awaiting the callback also invalidates the transfer stream; if + * the writer isn't done by then it'll receive an exception on the next write. + */ +template Fn> + requires requires(Fn fn, AsyncInputStream & i, AsyncIoRoot aio) { + [](kj::Promise>) {}(fn(i)); + } +auto wrapInAsyncPipe(Fn && fn) +{ + using promise_type = std::invoke_result_t; + using result_type = decltype(std::declval().wait(std::declval())); + + struct State + { + promise_type inner; + std::exception_ptr error; + bool finished = false; + + State(Fn && fn, std::unique_ptr reader) + : inner(fn(*reader) + .attach(std::move(reader)) + .then([this](auto result) { + if (result.has_error()) { + error = result.error(); + } + return result; + }) + .eagerlyEvaluate([this](kj::Exception && e) -> result_type { + try { + kj::throwFatalException(std::move(e)); + } catch (...) { + error = std::current_exception(); + throw; + } + })) + { + } + }; + + struct Writer : AsyncOutputStream + { + ref state; + std::unique_ptr out; + + Writer(ref state, std::unique_ptr out) : state(state), out(std::move(out)) {} + + kj::Promise> write(const void * src, size_t size) override + try { + if (state->error) { + co_return result::failure(state->error); + } else if (state->finished) { + throw Error("stream already closed"); + } + co_return LIX_TRY_AWAIT(out->write(src, size)); + } catch (...) { + // do not report this exception if `error` was set; we want to give + // priority to `runInner` for error reporting since the pipe itself + // is of little interest. if `inner` fails the pipe will break, all + // subsequent writes to the pipe will throw. only the `inner` error + // is actually interesting in this case, broken pipe errors aren't. + if (state->error) { + co_return result::failure(state->error); + } + state->error = std::current_exception(); + co_return result::current_exception(); + } + }; + + auto [reader, writer] = newZeroCopyPipe(); + auto state = make_ref(std::forward(fn), std::move(reader)); + return std::pair{ + std::make_unique(state, std::move(writer)), + kj::Function{[state]() { + state->finished = true; + return std::move(state->inner); + }}, + }; +} } diff --git a/tests/unit/libutil/async-io.cc b/tests/unit/libutil/async-io.cc index 53b7e9810..4b0969acf 100644 --- a/tests/unit/libutil/async-io.cc +++ b/tests/unit/libutil/async-io.cc @@ -273,4 +273,57 @@ TEST(AsyncZeroCopyPipe, dropWriter) (void) auto(std::move(w)); ASSERT_EQ(rp.wait(aio.kj.waitScope).value(), std::nullopt); } + +TEST(wrapInAsyncPipe, good) +{ + AsyncIoRoot aio; + + auto fn = [](AsyncInputStream & in) { return in.drain(); }; + auto [out, finish] = wrapInAsyncPipe(fn); + out->write("a", 1).wait(aio.kj.waitScope).value(); + out->write("b", 1).wait(aio.kj.waitScope).value(); + (void) auto(std::move(out)); + ASSERT_EQ(finish().wait(aio.kj.waitScope).value(), "ab"); +} + +TEST(wrapInAsyncPipe, earlyReaderExit) +{ + AsyncIoRoot aio; + + auto fn = [](AsyncInputStream & in) { return kj::Promise>{{1}}; }; + auto [out, finish] = wrapInAsyncPipe(fn); + // should throw "broken pipe", but result is still good + ASSERT_THROW(out->write("a", 1).wait(aio.kj.waitScope).value(), Error); + ASSERT_EQ(finish().wait(aio.kj.waitScope).value(), 1); +} + +TEST(wrapInAsyncPipe, earlyFinish) +{ + AsyncIoRoot aio; + + auto fn = [](AsyncInputStream & in) { return kj::Promise>{kj::NEVER_DONE}; }; + auto [out, finish] = wrapInAsyncPipe(fn); + auto p = finish(); + ASSERT_THROW(out->write("a", 1).wait(aio.kj.waitScope).value(), Error); + ASSERT_FALSE(p.poll(aio.kj.waitScope)); +} + +TEST(wrapInAsyncPipe, readerError) +{ + AsyncIoRoot aio; + + auto fn = [](AsyncInputStream & in) -> kj::Promise> { + try { + char buf; + TRY_AWAIT(in.read(&buf, 1)); + throw UnimplementedError("marker"); + } catch (...) { + co_return result::current_exception(); + } + }; + auto [out, finish] = wrapInAsyncPipe(fn); + out->write("a", 1).wait(aio.kj.waitScope).value(); + ASSERT_THROW(out->write("a", 1).wait(aio.kj.waitScope).value(), UnimplementedError); + ASSERT_THROW(finish().wait(aio.kj.waitScope).value(), UnimplementedError); +} }