libstore: use async streams in RemoteStore

this is a large step towards making RemoteStore a proper capnp rpc
interface, and it lets us get rid of the RemoteStore error handler
thread pool. this does mean we make six or more extra syscalls per
operation to set and clear socket non-blocking flags, but they are
pretty cheap compared to cross-thread wakeups and scheduling. once
we have real capnp rpc for store wires we can drop them again too.

Change-Id: I67dfebc8644a407cd4a8221ffcad02a938ac5abe
This commit is contained in:
eldritch horrors
2025-06-17 15:25:32 +02:00
parent 3f62905312
commit c3bc0d35dd
6 changed files with 151 additions and 157 deletions
+45 -41
View File
@@ -5,11 +5,13 @@
#include "lix/libstore/worker-protocol.hh"
#include "lix/libstore/worker-protocol-impl.hh"
#include "lix/libutil/async.hh"
#include "lix/libutil/async-io.hh"
#include "lix/libutil/file-descriptor.hh"
#include "lix/libutil/io-buffer.hh"
#include "lix/libutil/pool.hh"
#include "lix/libutil/result.hh"
#include "lix/libutil/serialise.hh"
#include "lix/libutil/signals.hh"
#include <kj/async.h>
#include <tuple>
#include <type_traits>
@@ -98,7 +100,7 @@ struct RemoteStore::Connection
std::exception_ptr e;
};
kj::Promise<Result<RemoteError>> processStderr();
kj::Promise<Result<RemoteError>> processStderr(AsyncFdIoStream & stream);
};
/**
@@ -111,82 +113,89 @@ struct RemoteStore::Connection
struct RemoteStore::ConnectionHandle
{
Pool<RemoteStore::Connection>::Handle handle;
Sync<ThreadPool> & handlerThreads;
ConnectionHandle(
Pool<RemoteStore::Connection>::Handle && handle, Sync<ThreadPool> & handlerThreads
)
: handle(std::move(handle))
, handlerThreads(handlerThreads)
{
}
ConnectionHandle(ConnectionHandle && h)
: handle(std::move(h.handle))
, handlerThreads(h.handlerThreads)
{
}
ConnectionHandle(Pool<RemoteStore::Connection>::Handle && handle) : handle(std::move(handle)) {}
RemoteStore::Connection & operator * () { return *handle; }
RemoteStore::Connection * operator -> () { return &*handle; }
kj::Promise<Result<void>> processStderr();
kj::Promise<Result<void>> processStderr(AsyncFdIoStream & stream);
kj::Promise<Result<void>>
withFramedSinkAsync(std::function<kj::Promise<Result<void>>(Sink & sink)> fun);
kj::Promise<Result<void>> withFramedStream(
AsyncFdIoStream & stream,
std::function<kj::Promise<Result<void>>(AsyncOutputStream & stream)> fun
);
template<typename R = void, typename... Args>
kj::Promise<Result<R>> sendCommand(Args &&... args)
kj::Promise<Result<R>> sendCommandUninterruptible(Args &&... args)
try {
constexpr auto LastArgIdx = sizeof...(Args) - 1;
using AllArgsT = std::tuple<Args &&...>;
using LastArgT = std::tuple_element_t<LastArgIdx, AllArgsT>;
// invalidate this connection if we're cancelled early, e.g. by a user ^C.
// regular exceptions must be handled elsewhere due the subframe requests.
// this also invalidates connections if a request was sent while unwinding
// the stack, but that's sufficiently suspect to warrant being as careful.
auto invalidateOnCancel = kj::defer([&] {
if (std::uncaught_exceptions() == 0) {
handle.markBad();
}
});
// 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(StringSink s) { s << std::declval<LastArgT>(); }) {
AsyncFdIoStream stream{AsyncFdIoStream::shared_fd{}, handle->getFD()};
try {
StringSink msg;
((msg << std::forward<Args>(args)), ...);
writeFull(handle->getFD(), msg.s);
TRY_AWAIT(stream.writeFull(msg.s.data(), msg.s.size()));
} catch (...) {
handle.markBad();
throw;
}
LIX_TRY_AWAIT(processStderr());
LIX_TRY_AWAIT(processStderr(stream));
} 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...>) {
AsyncFdIoStream stream{AsyncFdIoStream::shared_fd{}, handle->getFD()};
try {
StringSink msg;
((msg << std::forward<std::tuple_element_t<Ids, AllArgsT>>(
std::get<Ids>(allArgs)
)),
...);
writeFull(handle->getFD(), msg.s);
[&]<size_t... Ids>(std::integer_sequence<size_t, Ids...>) {
((msg << std::forward<std::tuple_element_t<Ids, AllArgsT>>(
std::get<Ids>(allArgs)
)),
...);
}(ImmediateArgsIdxs{});
TRY_AWAIT(stream.writeFull(msg.s.data(), msg.s.size()));
} catch (...) {
handle.markBad();
throw;
}
}(ImmediateArgsIdxs{});
LIX_TRY_AWAIT(withFramedSinkAsync(std::get<LastArgIdx>(allArgs)));
LIX_TRY_AWAIT(withFramedStream(stream, std::get<LastArgIdx>(allArgs)));
}
if constexpr (std::is_void_v<R>) {
invalidateOnCancel.cancel();
co_return result::success();
} else {
try {
// NOTE while no async streams are using the fd it is fully synchronous.
// we need either sync sources or async sources, and async sources would
// require a *lot* of code duplication. response messages are mostly not
// large enough to block us for long, so we just accept the hit for now.
FdSource from{handle->getFD(), handle->fromBuf};
from.specialEndOfFileError = "Nix daemon disconnected while reading a response";
co_return WorkerProto::Serialise<R>::read(
{from, *handle->store, handle->daemonVersion}
);
auto result =
WorkerProto::Serialise<R>::read({from, *handle->store, handle->daemonVersion});
invalidateOnCancel.cancel();
co_return result;
} catch (...) {
handle.markBad();
throw;
@@ -196,16 +205,11 @@ struct RemoteStore::ConnectionHandle
co_return result::current_exception();
}
private:
struct FramedSinkHandler
template<typename R = void, typename... Args>
kj::Promise<Result<R>> sendCommand(Args &&... args)
{
std::exception_ptr ex;
std::packaged_task<void(AsyncIoRoot &)> stderrHandler;
explicit FramedSinkHandler(ConnectionHandle & conn, ThreadPool & handlerThreads);
~FramedSinkHandler() noexcept(false);
};
return makeInterruptible(sendCommandUninterruptible<R>(std::forward<Args>(args)...));
}
};
}
+70 -59
View File
@@ -1,5 +1,7 @@
#include "lix/libutil/async-collect.hh"
#include "lix/libutil/async-io.hh"
#include "lix/libutil/async.hh"
#include "lix/libutil/box_ptr.hh"
#include "lix/libutil/error.hh"
#include "lix/libutil/result.hh"
#include "lix/libutil/serialise.hh"
@@ -20,13 +22,14 @@
#include "lix/libutil/logging.hh"
#include "lix/libstore/filetransfer.hh"
#include "lix/libutil/strings.hh"
#include "lix/libutil/thread-name.hh"
#include "lix/libutil/thread-pool.hh"
#include "lix/libutil/types.hh"
#include "path-info.hh"
#include <cstdint>
#include <kj/async.h>
#include <kj/common.h>
#include <optional>
#include <string_view>
#include <utility>
namespace nix {
@@ -78,6 +81,9 @@ try {
kj::Promise<Result<void>> RemoteStore::initConnection(Connection & conn)
try {
/* Send the magic greeting, check for the reply. */
// NOTE: this is synchronous until we call processStderr. this is intentional;
// reading the response would be synchronous anyway, and making sending of the
// greeting synchronous also makes this code path significantly more readable.
try {
conn.store = this;
FdSource from{conn.getFD(), conn.fromBuf};
@@ -109,7 +115,8 @@ try {
{from, *conn.store, conn.daemonVersion}
);
auto ex = TRY_AWAIT(conn.processStderr());
AsyncFdIoStream stream{AsyncFdIoStream::shared_fd{}, conn.getFD()};
auto ex = TRY_AWAIT(conn.processStderr(stream));
if (ex.e) {
std::rethrow_exception(ex.e);
}
@@ -161,8 +168,9 @@ try {
for (auto & i : overrides)
command << i.first << i.second.value;
writeFull(conn.getFD(), command.s);
auto ex = TRY_AWAIT(conn.processStderr());
AsyncFdIoStream stream{AsyncFdIoStream::shared_fd{}, conn.getFD()};
TRY_AWAIT(stream.writeFull(command.s.data(), command.s.size()));
auto ex = TRY_AWAIT(conn.processStderr(stream));
if (ex.e) {
std::rethrow_exception(ex.e);
}
@@ -171,9 +179,9 @@ try {
co_return result::current_exception();
}
kj::Promise<Result<void>> RemoteStore::ConnectionHandle::processStderr()
kj::Promise<Result<void>> RemoteStore::ConnectionHandle::processStderr(AsyncFdIoStream & stream)
try {
auto ex = TRY_AWAIT(handle->processStderr());
auto ex = TRY_AWAIT(handle->processStderr(stream));
if (ex.e) {
co_return result::failure(ex.e);
}
@@ -185,7 +193,7 @@ try {
kj::Promise<Result<RemoteStore::ConnectionHandle>> RemoteStore::getConnection()
try {
co_return ConnectionHandle(TRY_AWAIT(connections->get()), handlerThreads);
co_return ConnectionHandle(TRY_AWAIT(connections->get()));
} catch (...) {
co_return result::current_exception();
}
@@ -372,7 +380,7 @@ try {
caMethod.render(hashType),
WorkerProto::write(*conn, references),
repair,
[&](Sink & sink) { return dump.drainInto(sink); }
[&](AsyncOutputStream & stream) { return dump.drainInto(stream); }
)));
} catch (...) {
co_return result::current_exception();
@@ -417,7 +425,7 @@ try {
renderContentAddress(info.ca),
repair,
!checkSigs,
[&](Sink & sink) { return copier->drainInto(sink); }
[&](AsyncOutputStream & stream) { return copier->drainInto(stream); }
));
co_return result::success();
} catch (...) {
@@ -439,14 +447,20 @@ try {
repair,
!checkSigs,
// NOLINTNEXTLINE(cppcoreguidelines-avoid-capturing-lambda-coroutines)
[&](Sink & sink) -> kj::Promise<Result<void>> {
[&](AsyncOutputStream & stream) -> kj::Promise<Result<void>> {
try {
sink << pathsToCopy.size();
auto send = [&]<typename T>(T && value) {
auto tmp = make_box_ptr<StringSink>();
*tmp << std::forward<T>(value);
return stream.writeFull(tmp->s.data(), tmp->s.size()).attach(std::move(tmp));
};
TRY_AWAIT(send(pathsToCopy.size()));
for (auto & [pathInfo, pathSource] : pathsToCopy) {
sink << WorkerProto::Serialise<ValidPathInfo>::write(
TRY_AWAIT(send(WorkerProto::Serialise<ValidPathInfo>::write(
WorkerProto::WriteConn{*this, remoteVersion}, pathInfo
);
TRY_AWAIT(TRY_AWAIT(pathSource())->drainInto(sink));
)));
TRY_AWAIT(TRY_AWAIT(pathSource())->drainInto(stream));
}
co_return result::success();
} catch (...) {
@@ -664,7 +678,7 @@ try {
TRY_AWAIT(conn.sendCommand<unsigned>(
WorkerProto::Op::AddBuildLog,
drvPath.to_string(),
[&](Sink & sink) { return source.drainInto(sink); }
[&](AsyncOutputStream & stream) { return source.drainInto(stream); }
));
co_return result::success();
} catch (...) {
@@ -757,9 +771,28 @@ static Logger::Fields readFields(Source & from)
return fields;
}
kj::Promise<Result<RemoteStore::Connection::RemoteError>> RemoteStore::Connection::processStderr()
kj::Promise<Result<RemoteStore::Connection::RemoteError>>
RemoteStore::Connection::processStderr(AsyncFdIoStream & stream)
try {
while (true) {
// fill the read buffer asynchronously until we have at least a message type.
// once we have this we'll continue synchronously as in the sendCommand case.
while (fromBuf->used() < sizeof(uint64_t)) {
const auto available = fromBuf->getWriteBuffer();
fromBuf->added(TRY_AWAIT(stream.read(available.data(), available.size())));
}
// SAFETY NOTE: while we're running we own the executor, and thus the stream.
// setting these flags is unsafe if the stream is shared with another thread.
const int oldFlags = fcntl(getFD(), F_GETFL);
if (oldFlags == -1 || fcntl(getFD(), F_SETFL, oldFlags & ~O_NONBLOCK) < 0) {
throw SysError("making connection blocking");
}
KJ_DEFER({
if (fcntl(getFD(), F_SETFL, oldFlags) < 0) {
throw SysError("restoring connection flags");
}
});
FdSource from{getFD(), fromBuf};
from.specialEndOfFileError = "Nix daemon disconnected while waiting for a response";
@@ -807,52 +840,30 @@ try {
co_return result::current_exception();
}
RemoteStore::ConnectionHandle::FramedSinkHandler::FramedSinkHandler(
ConnectionHandle & conn, ThreadPool & handlerThreads
)
: stderrHandler([&](AsyncIoRoot & aio) {
try {
aio.blockOn(conn.processStderr());
} catch (...) {
ex = std::current_exception();
}
})
{
handlerThreads.enqueueWithAio([&](AsyncIoRoot & aio) { stderrHandler(aio); });
}
RemoteStore::ConnectionHandle::FramedSinkHandler::~FramedSinkHandler() noexcept(false)
{
stderrHandler.get_future().get();
// if we're handling an Interrupted exception we must be careful: it's
// possible that the exception was thrown by the withFramedSink framed
// function, but not by the FramedSink itself. in this case our stderr
// handler thread may race with FramedSink::writeUnbuffered, catch the
// Interrupted exception independently, store it into ex, and have our
// own destructor rethrow a second copy of Interrupted. since we can't
// handle multiple exceptions anyway the safest path is to simply drop
// the remote (possibly Interrupted) exception when called for unwind.
if (ex && std::uncaught_exceptions() == 0) {
throw FramedSink::RemoteError(ex);
}
}
kj::Promise<Result<void>> RemoteStore::ConnectionHandle::withFramedSinkAsync(
std::function<kj::Promise<Result<void>>(Sink & sink)> fun
kj::Promise<Result<void>> RemoteStore::ConnectionHandle::withFramedStream(
AsyncFdIoStream & stream,
std::function<kj::Promise<Result<void>>(AsyncOutputStream & stream)> fun
)
try {
{
FdSink to{handle->getFD()};
FramedSinkHandler handler{*this, *handlerThreads.lock()};
FramedSink sink(to, handler.ex);
TRY_AWAIT(fun(sink));
sink.flush();
}
AsyncBufferedOutputStream to(stream);
AsyncFramedStream sink(to);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-capturing-lambda-coroutines)
auto send = [&]() -> kj::Promise<Result<void>> {
try {
TRY_AWAIT(fun(sink));
TRY_AWAIT(sink.finish());
TRY_AWAIT(to.flush());
co_return result::success();
} catch (...) {
handle.markBad();
co_return result::current_exception();
}
};
TRY_AWAIT(asyncJoin(send(), processStderr(stream)));
co_return result::success();
} catch (FramedSink::RemoteError & e) {
co_return result::failure(e.e);
} catch (...) {
handle.markBad();
co_return result::current_exception();
}
}
-8
View File
@@ -9,7 +9,6 @@
#include "lix/libstore/gc-store.hh"
#include "lix/libstore/log-store.hh"
#include "lix/libutil/async-io.hh"
#include "lix/libutil/thread-pool.hh"
#include "lix/libutil/types.hh"
@@ -204,13 +203,6 @@ private:
std::atomic_bool failed{false};
// NOTE we rely on the thread pool not starting threads eagerly. if it ever starts
// doing that we're certainly going to fail due to the immense thread count, which
// we need to satisfy temporary `incCapacity` calls by some RemoteStore functions.
Sync<ThreadPool> handlerThreads{
std::in_place, "remote stderr", std::numeric_limits<size_t>::max()
};
kj::Promise<Result<void>> copyDrvsFromEvalStore(
const std::vector<DerivedPath> & paths,
std::shared_ptr<Store> evalStore);
+21
View File
@@ -190,4 +190,25 @@ kj::Promise<Result<size_t>> AsyncFdIoStream::write(const void * src, size_t size
return {result::failure(std::make_exception_ptr(SysError(errno, "write failed")))};
}
}
kj::Promise<Result<void>> AsyncFramedStream::finish()
try {
StringSink tmp;
tmp << 0;
TRY_AWAIT(to.writeFull(tmp.s.data(), tmp.s.size()));
co_return result::success();
} catch (...) {
co_return result::current_exception();
}
kj::Promise<Result<size_t>> AsyncFramedStream::write(const void * buffer, size_t size)
try {
StringSink tmp;
tmp << size;
TRY_AWAIT(to.writeFull(tmp.s.data(), tmp.s.size()));
TRY_AWAIT(to.writeFull(buffer, size));
co_return size;
} catch (...) {
co_return result::current_exception();
}
}
+15
View File
@@ -182,4 +182,19 @@ public:
kj::Promise<Result<size_t>> read(void * tgt, size_t size) override;
kj::Promise<Result<size_t>> write(const void * src, size_t size) override;
};
/**
* Write as chunks in the format expected by FramedSource.
*/
class AsyncFramedStream : public AsyncOutputStream
{
AsyncOutputStream & to;
public:
explicit AsyncFramedStream(AsyncOutputStream & to) : to(to) {}
kj::Promise<Result<void>> finish();
kj::Promise<Result<size_t>> write(const void * src, size_t size) override;
};
}
-49
View File
@@ -544,53 +544,4 @@ struct FramedSource : Source
return n;
}
};
/**
* Write as chunks in the format expected by FramedSource.
*
* The exception_ptr reference can be used to terminate the stream when you
* detect that an error has occurred on the remote end.
*/
struct FramedSink : nix::BufferedSink
{
BufferedSink & to;
std::exception_ptr & ex;
struct RemoteError : BaseException
{
std::exception_ptr e;
RemoteError(std::exception_ptr e)
: e(e) // NOLINT(bugprone-throw-keyword-missing): intentional copy
{
}
};
FramedSink(BufferedSink & to, std::exception_ptr & ex) : to(to), ex(ex)
{ }
~FramedSink()
{
try {
to << 0;
to.flush();
} catch (...) {
ignoreExceptionInDestructor();
}
}
void writeUnbuffered(std::string_view data) override
{
/* Don't send more data if the remote has
encountered an error. */
if (ex) {
auto ex2 = ex;
ex = nullptr;
throw RemoteError{ex2};
}
to << data.size();
to(data);
};
};
}