diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index fd930d4e..38909ccb 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -1,5 +1,5 @@ { - "lockFileVersion": 26, + "lockFileVersion": 28, "registryFileHashes": { "https://bcr.bazel.build/bazel_registry.json": "8a28e4aff06ee60aed2a8c281907fb8bcbf3b753c91fb5a5c57da3215d5b3497", "https://bcr.bazel.build/modules/abseil-cpp/20210324.2/MODULE.bazel": "7cd0312e064fde87c8d1cd79ba06c876bd23630c83466e9500321be55c96ace2", @@ -324,22 +324,9 @@ }, "selectedYankedVersions": {}, "moduleExtensions": { - "@@rules_android_ndk+//:extension.bzl%android_ndk_repository_extension": { - "general": { - "bzlTransitiveDigest": "j+NMXk5/a1vnDypModKL+ToEo/4o7DJeXSFhTfbJ7rI=", - "usagesDigest": "Dshe+RJHb2CHJNFJG53h45menXtSyF3M1AMuHpTXiGU=", - "recordedInputs": [], - "generatedRepoSpecs": { - "androidndk": { - "repoRuleId": "@@rules_android_ndk+//:rules.bzl%android_ndk_repository", - "attributes": {} - } - } - } - }, "@@rules_kotlin+//src/main/starlark/core/repositories:bzlmod_setup.bzl%rules_kotlin_extensions": { "general": { - "bzlTransitiveDigest": "Ga4z8lQy1YQ5rAMy+dOl0dqcCEBnYNCXku8x3YQmDZI=", + "bzlTransitiveDigest": "+Kp6j204mBZ3mxlIDDR0gBoP45BZ4jYRhRAcB8sU0qc=", "usagesDigest": "QI2z8ZUR+mqtbwsf2fLqYdJAkPOHdOV+tF2yVAUgRzw=", "recordedInputs": [ "REPO_MAPPING:rules_kotlin+,bazel_tools bazel_tools" @@ -396,7 +383,7 @@ }, "@@rules_python+//python/extensions:config.bzl%config": { "general": { - "bzlTransitiveDigest": "iibnRYgg8LpcfmH7EAnVwYePC3jsVaJ6Id8XxUjSZps=", + "bzlTransitiveDigest": "dzD8Q2YmrP3fz8saWLHPmlwPLO91ImtTmP/c9JKTStM=", "usagesDigest": "ZVSXMAGpD+xzVNPuvF1IoLBkty7TROO0+akMapt1pAg=", "recordedInputs": [ "REPO_MAPPING:rules_python+,bazel_tools bazel_tools", @@ -594,8 +581,8 @@ }, "@@rules_rust+//crate_universe:extensions.bzl%crate": { "general": { - "bzlTransitiveDigest": "qD4RshABgOGJEdLYZOfeternLk7PjIxk+jQHfa/mw+Q=", - "usagesDigest": "ae+fUjKn/otf5JUNjlk1r6plMRIByMwUlo4LYuO1A1E=", + "bzlTransitiveDigest": "eT6cMKC8/lOF0/pWxtwFM9dDgUj6wYzAHDNWS/fBCZ8=", + "usagesDigest": "4fRZ8PQG6B7QQECVGU7Eu39hXHHWXBRNnJqTZgXjkUY=", "recordedInputs": [ "ENV:CARGO_BAZEL_DEBUG \\0", "ENV:CARGO_BAZEL_GENERATOR_SHA256 \\0", @@ -1518,106 +1505,8 @@ } } } - }, - "@@rules_rust+//crate_universe/private:internal_extensions.bzl%cu_nr": { - "general": { - "bzlTransitiveDigest": "E+J8HoPZNy7ImtRKAd7uoNNklRekTbBcZIhLpcw0D2k=", - "usagesDigest": "v4We18mWSPeKV4GPp9Gne78W+jZOgP2pC1i4UN9br1g=", - "recordedInputs": [ - "REPO_MAPPING:bazel_features+,bazel_features_globals bazel_features++version_extension+bazel_features_globals", - "REPO_MAPPING:bazel_features+,bazel_features_version bazel_features++version_extension+bazel_features_version", - "REPO_MAPPING:rules_cc+,bazel_skylib bazel_skylib+", - "REPO_MAPPING:rules_cc+,bazel_tools bazel_tools", - "REPO_MAPPING:rules_cc+,cc_compatibility_proxy rules_cc++compatibility_proxy+cc_compatibility_proxy", - "REPO_MAPPING:rules_cc+,platforms platforms", - "REPO_MAPPING:rules_cc+,rules_cc rules_cc+", - "REPO_MAPPING:rules_cc++compatibility_proxy+cc_compatibility_proxy,rules_cc rules_cc+", - "REPO_MAPPING:rules_rust+,bazel_features bazel_features+", - "REPO_MAPPING:rules_rust+,bazel_skylib bazel_skylib+", - "REPO_MAPPING:rules_rust+,bazel_tools bazel_tools", - "REPO_MAPPING:rules_rust+,cui rules_rust++cu+cui", - "REPO_MAPPING:rules_rust+,rrc rules_rust++i2+rrc", - "REPO_MAPPING:rules_rust+,rules_cc rules_cc+", - "REPO_MAPPING:rules_rust+,rules_rust rules_rust+" - ], - "generatedRepoSpecs": { - "cargo_bazel_bootstrap": { - "repoRuleId": "@@rules_rust+//cargo/private:cargo_bootstrap.bzl%cargo_bootstrap_repository", - "attributes": { - "srcs": [ - "@@rules_rust+//crate_universe:src/api.rs", - "@@rules_rust+//crate_universe:src/api/lockfile.rs", - "@@rules_rust+//crate_universe:src/cli.rs", - "@@rules_rust+//crate_universe:src/cli/generate.rs", - "@@rules_rust+//crate_universe:src/cli/query.rs", - "@@rules_rust+//crate_universe:src/cli/render.rs", - "@@rules_rust+//crate_universe:src/cli/splice.rs", - "@@rules_rust+//crate_universe:src/cli/vendor.rs", - "@@rules_rust+//crate_universe:src/config.rs", - "@@rules_rust+//crate_universe:src/context.rs", - "@@rules_rust+//crate_universe:src/context/crate_context.rs", - "@@rules_rust+//crate_universe:src/context/platforms.rs", - "@@rules_rust+//crate_universe:src/lib.rs", - "@@rules_rust+//crate_universe:src/lockfile.rs", - "@@rules_rust+//crate_universe:src/main.rs", - "@@rules_rust+//crate_universe:src/metadata.rs", - "@@rules_rust+//crate_universe:src/metadata/cargo_bin.rs", - "@@rules_rust+//crate_universe:src/metadata/cargo_tree_resolver.rs", - "@@rules_rust+//crate_universe:src/metadata/cargo_tree_rustc_wrapper.bat", - "@@rules_rust+//crate_universe:src/metadata/cargo_tree_rustc_wrapper.sh", - "@@rules_rust+//crate_universe:src/metadata/dependency.rs", - "@@rules_rust+//crate_universe:src/metadata/metadata_annotation.rs", - "@@rules_rust+//crate_universe:src/rendering.rs", - "@@rules_rust+//crate_universe:src/rendering/template_engine.rs", - "@@rules_rust+//crate_universe:src/rendering/templates/module_bzl.j2", - "@@rules_rust+//crate_universe:src/rendering/templates/partials/header.j2", - "@@rules_rust+//crate_universe:src/rendering/templates/partials/module/aliases_map.j2", - "@@rules_rust+//crate_universe:src/rendering/templates/partials/module/deps_map.j2", - "@@rules_rust+//crate_universe:src/rendering/templates/partials/module/repo_git.j2", - "@@rules_rust+//crate_universe:src/rendering/templates/partials/module/repo_http.j2", - "@@rules_rust+//crate_universe:src/rendering/templates/vendor_module.j2", - "@@rules_rust+//crate_universe:src/rendering/verbatim/alias_rules.bzl", - "@@rules_rust+//crate_universe:src/select.rs", - "@@rules_rust+//crate_universe:src/splicing.rs", - "@@rules_rust+//crate_universe:src/splicing/cargo_config.rs", - "@@rules_rust+//crate_universe:src/splicing/crate_index_lookup.rs", - "@@rules_rust+//crate_universe:src/splicing/splicer.rs", - "@@rules_rust+//crate_universe:src/test.rs", - "@@rules_rust+//crate_universe:src/utils.rs", - "@@rules_rust+//crate_universe:src/utils/starlark.rs", - "@@rules_rust+//crate_universe:src/utils/starlark/glob.rs", - "@@rules_rust+//crate_universe:src/utils/starlark/label.rs", - "@@rules_rust+//crate_universe:src/utils/starlark/select.rs", - "@@rules_rust+//crate_universe:src/utils/starlark/select_dict.rs", - "@@rules_rust+//crate_universe:src/utils/starlark/select_list.rs", - "@@rules_rust+//crate_universe:src/utils/starlark/select_scalar.rs", - "@@rules_rust+//crate_universe:src/utils/starlark/select_set.rs", - "@@rules_rust+//crate_universe:src/utils/starlark/serialize.rs", - "@@rules_rust+//crate_universe:src/utils/starlark/target_compatible_with.rs", - "@@rules_rust+//crate_universe:src/utils/symlink.rs", - "@@rules_rust+//crate_universe:src/utils/target_triple.rs" - ], - "binary": "cargo-bazel", - "cargo_lockfile": "@@rules_rust+//crate_universe:Cargo.lock", - "cargo_toml": "@@rules_rust+//crate_universe:Cargo.toml", - "version": "1.86.0", - "timeout": 900, - "rust_toolchain_cargo_template": "@rust_host_tools//:bin/{tool}", - "rust_toolchain_rustc_template": "@rust_host_tools//:bin/{tool}", - "compressed_windows_toolchain_names": false - } - } - }, - "moduleExtensionMetadata": { - "explicitRootModuleDirectDeps": [ - "cargo_bazel_bootstrap" - ], - "explicitRootModuleDirectDevDeps": [], - "useAllRepos": "NO", - "reproducible": false - } - } } }, - "facts": {} + "facts": {}, + "factsVersions": {} } diff --git a/asio_rpc/server/rpc_server.cc b/asio_rpc/server/rpc_server.cc index e0761a3f..67a0121f 100644 --- a/asio_rpc/server/rpc_server.cc +++ b/asio_rpc/server/rpc_server.cc @@ -8,6 +8,7 @@ #include "proto/subspace.pb.h" #include #include +#include namespace subspace::asio_rpc { @@ -763,27 +764,36 @@ void RpcServer::SessionStreamingMethodCoroutine( method_instance->method->name.c_str()); AnyStreamWriter writer(server, session, method_instance, request); + auto stop_pipe_or = toolbelt::Pipe::Create(); + if (!stop_pipe_or.ok()) { + server->logger_.Log(toolbelt::LogLevel::kError, + "Failed to create stream cancellation stop pipe: %s", + stop_pipe_or.status().ToString().c_str()); + continue; + } + auto stop_pipe = + std::make_shared(std::move(*stop_pipe_or)); + auto writer_state = writer.state; + const int32_t session_id = session->session_id; + const int32_t request_id = request.request_id(); // Spawn a coroutine to read the cancellation channel. boost::asio::spawn( *server->io_context_, - [server, session, method_instance, &writer, - &request](boost::asio::yield_context yield) { - int dup_interrupt = - ::dup(server->interrupt_pipe_.ReadFd().Fd()); - toolbelt::FileDescriptor interrupt(dup_interrupt); - while (!writer.IsCancelled()) { + [server, method_instance, stop_pipe, writer_state, session_id, + request_id](boost::asio::yield_context yield) { + while (!writer_state->is_cancelled.load(std::memory_order_acquire)) { int cancel_fd = method_instance->cancel_subscriber->GetPollFd().fd; auto s = async_wait_either(*server->io_context_, cancel_fd, - interrupt.Fd(), yield); + stop_pipe->ReadFd().Fd(), yield); if (!s.ok()) { server->logger_.Log(toolbelt::LogLevel::kError, "Error waiting for cancel: %s", s.status().ToString().c_str()); return; } - if (*s == interrupt.Fd()) { + if (*s == stop_pipe->ReadFd().Fd()) { break; } bool cancel_ok = false; @@ -805,14 +815,14 @@ void RpcServer::SessionStreamingMethodCoroutine( msg.status().ToString().c_str()); continue; } - if (cancel.session_id() == session->session_id && - cancel.request_id() == request.request_id()) { + if (cancel.session_id() == session_id && + cancel.request_id() == request_id) { cancel_ok = true; break; } } if (cancel_ok) { - writer.Cancel(); + writer_state->is_cancelled.store(true, std::memory_order_release); } } }, @@ -832,6 +842,8 @@ void RpcServer::SessionStreamingMethodCoroutine( method_status.ToString()), yield); } + char stop = 1; + (void)::write(stop_pipe->WriteFd().Fd(), &stop, 1); } } diff --git a/asio_rpc/server/rpc_server.h b/asio_rpc/server/rpc_server.h index 7205c890..d2366a05 100644 --- a/asio_rpc/server/rpc_server.h +++ b/asio_rpc/server/rpc_server.h @@ -14,6 +14,7 @@ #include "toolbelt/pipe.h" #include +#include #include namespace subspace::asio_rpc { @@ -25,27 +26,34 @@ namespace internal { struct Session; struct MethodInstance; +struct StreamWriterState { + std::atomic_bool is_cancelled{false}; +}; + struct AnyStreamWriter { AnyStreamWriter(std::shared_ptr server, std::shared_ptr session, std::shared_ptr method_instance, const RpcRequest &request) : server(std::move(server)), session(std::move(session)), - method_instance(std::move(method_instance)), request(request) {} + method_instance(std::move(method_instance)), request(request), + state(std::make_shared()) {} bool Write(std::unique_ptr res, boost::asio::yield_context yield); void Finish(boost::asio::yield_context yield); - void Cancel() { is_cancelled = true; } + void Cancel() { state->is_cancelled.store(true, std::memory_order_release); } - bool IsCancelled() const { return is_cancelled; } + bool IsCancelled() const { + return state->is_cancelled.load(std::memory_order_acquire); + } std::shared_ptr server; std::shared_ptr session; std::shared_ptr method_instance; - const RpcRequest &request; - bool is_cancelled = false; + RpcRequest request; + std::shared_ptr state; }; // An item pushed by either the reply function or the error function of an diff --git a/asio_rpc/server/server_test.cc b/asio_rpc/server/server_test.cc index c33e97d7..c812276f 100644 --- a/asio_rpc/server/server_test.cc +++ b/asio_rpc/server/server_test.cc @@ -133,6 +133,18 @@ static std::shared_ptr BuildServer() { return server; } +TEST_F(AsioServerTest, StreamWriterOwnsRequestMetadata) { + std::unique_ptr writer; + { + subspace::RpcRequest request; + request.set_request_id(1234); + writer = std::make_unique( + nullptr, nullptr, nullptr, request); + } + + EXPECT_EQ(1234, writer->request.request_id()); +} + struct ServerContext { std::shared_ptr client; std::shared_ptr pub; diff --git a/client/client.cc b/client/client.cc index 631446ed..2df7bcb7 100644 --- a/client/client.cc +++ b/client/client.cc @@ -885,9 +885,11 @@ ClientImpl::WaitForReliablePublisher(PublisherImpl *publisher, uint64_t timeout_ns = timeout.count(); #if SUBSPACE_CORO_BACKEND == SUBSPACE_CORO_BACKEND_ASIO if (IsCooperative()) { - // The Asio WaitEither has no timeout; it waits until one fd is ready. absl::StatusOr r = async::WaitEither( - SocketContext(), publisher->GetPollFd().Fd(), fd.Fd()); + SocketContext(), publisher->GetPollFd().Fd(), fd.Fd(), timeout); + if (absl::IsDeadlineExceeded(r.status())) { + return absl::InternalError("Timeout waiting for reliable publisher"); + } if (!r.ok()) { return r.status(); } @@ -984,9 +986,11 @@ absl::StatusOr ClientImpl::WaitForSubscriber( uint64_t timeout_ns = timeout.count(); #if SUBSPACE_CORO_BACKEND == SUBSPACE_CORO_BACKEND_ASIO if (IsCooperative()) { - // The Asio WaitEither has no timeout; it waits until one fd is ready. absl::StatusOr r = async::WaitEither( - SocketContext(), subscriber->GetPollFd().Fd(), fd.Fd()); + SocketContext(), subscriber->GetPollFd().Fd(), fd.Fd(), timeout); + if (absl::IsDeadlineExceeded(r.status())) { + return absl::InternalError("Timeout waiting for subscriber"); + } if (!r.ok()) { return r.status(); } @@ -1073,8 +1077,12 @@ ClientImpl::WaitForReliablePublisher(PublisherImpl *publisher, !status.ok()) { return status; } - // The Asio WaitEither has no timeout; it waits until one fd is ready. - return async::WaitEither(ctx, publisher->GetPollFd().Fd(), fd.Fd()); + absl::StatusOr r = + async::WaitEither(ctx, publisher->GetPollFd().Fd(), fd.Fd(), timeout); + if (absl::IsDeadlineExceeded(r.status())) { + return absl::InternalError("Timeout waiting for reliable publisher"); + } + return r; } absl::Status ClientImpl::WaitForSubscriber(SubscriberImpl *subscriber, @@ -1097,8 +1105,12 @@ absl::StatusOr ClientImpl::WaitForSubscriber( if (absl::Status status = CheckConnected(); !status.ok()) { return status; } - // The Asio WaitEither has no timeout; it waits until one fd is ready. - return async::WaitEither(ctx, subscriber->GetPollFd().Fd(), fd.Fd()); + absl::StatusOr r = + async::WaitEither(ctx, subscriber->GetPollFd().Fd(), fd.Fd(), timeout); + if (absl::IsDeadlineExceeded(r.status())) { + return absl::InternalError("Timeout waiting for subscriber"); + } + return r; } #endif diff --git a/co20_rpc/server/rpc_server.cc b/co20_rpc/server/rpc_server.cc index 57e3f633..4d0f8774 100644 --- a/co20_rpc/server/rpc_server.cc +++ b/co20_rpc/server/rpc_server.cc @@ -7,6 +7,7 @@ #include "proto/subspace.pb.h" #include #include +#include namespace subspace::co20_rpc { @@ -738,12 +739,24 @@ co20::ValueTask RpcServer::SessionStreamingMethodCoroutine( method_instance->method->name.c_str()); AnyStreamWriter writer(server, session, method_instance, request); + auto stop_pipe_or = toolbelt::Pipe::Create(); + if (!stop_pipe_or.ok()) { + server->logger_.Log(toolbelt::LogLevel::kError, + "Failed to create stream cancellation stop pipe: %s", + stop_pipe_or.status().ToString().c_str()); + continue; + } + auto stop_pipe = + std::make_shared(std::move(*stop_pipe_or)); + auto writer_state = writer.state; + const int32_t session_id = session->session_id; + const int32_t request_id = request.request_id(); // Spawn a coroutine to read the cancellation channel. server->scheduler_->Spawn( - [server, session, method_instance, &writer, - &request](co20::Coroutine &cancel_co) -> co20::Task { - while (!writer.IsCancelled()) { + [server, method_instance, stop_pipe, writer_state, session_id, + request_id](co20::Coroutine &cancel_co) -> co20::Task { + while (!writer_state->is_cancelled.load(std::memory_order_acquire)) { int cancel_fd = method_instance->cancel_subscriber->GetPollFd().fd; int fd = co_await cancel_co.Wait(cancel_fd, POLLIN); @@ -769,19 +782,19 @@ co20::ValueTask RpcServer::SessionStreamingMethodCoroutine( msg.status().ToString().c_str()); continue; } - if (cancel.session_id() == session->session_id && - cancel.request_id() == request.request_id()) { + if (cancel.session_id() == session_id && + cancel.request_id() == request_id) { cancel_ok = true; break; } } if (cancel_ok) { - writer.Cancel(); + writer_state->is_cancelled.store(true, std::memory_order_release); } } co_return; }, - "cancel_watcher", server->interrupt_pipe_.ReadFd().Fd()); + "cancel_watcher", stop_pipe->ReadFd().Fd()); absl::Status method_status = co_await method_instance->method->stream_callback(request.argument(), @@ -797,6 +810,8 @@ co20::ValueTask RpcServer::SessionStreamingMethodCoroutine( method_instance->method->name, method_status.ToString())); } + char stop = 1; + (void)::write(stop_pipe->WriteFd().Fd(), &stop, 1); } co_return; } diff --git a/co20_rpc/server/rpc_server.h b/co20_rpc/server/rpc_server.h index 8bdf028e..67ac7da0 100644 --- a/co20_rpc/server/rpc_server.h +++ b/co20_rpc/server/rpc_server.h @@ -14,6 +14,8 @@ #include "toolbelt/logging.h" #include "toolbelt/pipe.h" +#include + namespace subspace::co20_rpc { class RpcServer; @@ -23,26 +25,33 @@ namespace internal { struct Session; struct MethodInstance; +struct StreamWriterState { + std::atomic_bool is_cancelled{false}; +}; + struct AnyStreamWriter { AnyStreamWriter(std::shared_ptr server, std::shared_ptr session, std::shared_ptr method_instance, const RpcRequest &request) : server(std::move(server)), session(std::move(session)), - method_instance(std::move(method_instance)), request(request) {} + method_instance(std::move(method_instance)), request(request), + state(std::make_shared()) {} co20::ValueTask Write(std::unique_ptr res); co20::ValueTask Finish(); - void Cancel() { is_cancelled = true; } + void Cancel() { state->is_cancelled.store(true, std::memory_order_release); } - bool IsCancelled() const { return is_cancelled; } + bool IsCancelled() const { + return state->is_cancelled.load(std::memory_order_acquire); + } std::shared_ptr server; std::shared_ptr session; std::shared_ptr method_instance; - const RpcRequest &request; - bool is_cancelled = false; + RpcRequest request; + std::shared_ptr state; }; // An item pushed by either the reply function or the error function of an diff --git a/co20_rpc/server/server_test.cc b/co20_rpc/server/server_test.cc index b500733b..aeea68bc 100644 --- a/co20_rpc/server/server_test.cc +++ b/co20_rpc/server/server_test.cc @@ -135,6 +135,18 @@ static std::shared_ptr BuildServer() { return server; } +TEST_F(Co20ServerTest, StreamWriterOwnsRequestMetadata) { + std::unique_ptr writer; + { + subspace::RpcRequest request; + request.set_request_id(1234); + writer = std::make_unique( + nullptr, nullptr, nullptr, request); + } + + EXPECT_EQ(1234, writer->request.request_id()); +} + struct ServerContext { std::shared_ptr client; std::shared_ptr pub; diff --git a/common/async/async_test.cc b/common/async/async_test.cc index 2660fd44..8cd6899c 100644 --- a/common/async/async_test.cc +++ b/common/async/async_test.cc @@ -12,7 +12,9 @@ #include #if SUBSPACE_CORO_BACKEND == SUBSPACE_CORO_BACKEND_ASIO +#include #include +#include #endif namespace subspace::async { @@ -78,6 +80,40 @@ TEST(AsyncTest, WaitReadableTimeout) { ::close(fds[1]); } +#if SUBSPACE_CORO_BACKEND == SUBSPACE_CORO_BACKEND_ASIO +TEST(AsyncTest, WaitReadableWorksWithoutCancellationSlot) { + int fds[2]; + ASSERT_EQ(0, ::pipe(fds)); + + boost::asio::io_context ioc; + bool got_data = false; + + boost::asio::spawn( + ioc, + [&](Context ctx) { + ASSERT_TRUE( + WaitReadable(ctx, fds[0], std::chrono::milliseconds(100)).ok()); + char buf[8] = {}; + ssize_t n = ::read(fds[0], buf, sizeof(buf)); + got_data = (n == 1 && buf[0] == 'x'); + }, + boost::asio::detached); + + boost::asio::spawn( + ioc, + [&](Context ctx) { + Sleep(ctx, std::chrono::milliseconds(10)); + ASSERT_EQ(1, ::write(fds[1], "x", 1)); + }, + boost::asio::detached); + + ioc.run(); + EXPECT_TRUE(got_data); + ::close(fds[0]); + ::close(fds[1]); +} +#endif + TEST(AsyncTest, WaitEitherReturnsReadyFd) { int a[2], b[2]; ASSERT_EQ(0, ::pipe(a)); @@ -106,6 +142,28 @@ TEST(AsyncTest, WaitEitherReturnsReadyFd) { ::close(b[1]); } +TEST(AsyncTest, WaitEitherTimeout) { + int a[2], b[2]; + ASSERT_EQ(0, ::pipe(a)); + ASSERT_EQ(0, ::pipe(b)); + + RuntimeFixture fx; + bool timed_out = false; + + fx.runtime.Spawn([&](Context ctx) { + absl::StatusOr r = + WaitEither(ctx, a[0], b[0], std::chrono::milliseconds(20)); + timed_out = absl::IsDeadlineExceeded(r.status()); + }); + + fx.Run(); + EXPECT_TRUE(timed_out); + ::close(a[0]); + ::close(a[1]); + ::close(b[0]); + ::close(b[1]); +} + // A round-trip over the stream socket facade, exercising the length-delimited // framing on the loopback interface. TEST(AsyncTest, StreamSocketLoopbackMessage) { diff --git a/common/async/runtime.h b/common/async/runtime.h index ae0b01fd..00a5d227 100644 --- a/common/async/runtime.h +++ b/common/async/runtime.h @@ -111,6 +111,22 @@ class AsyncRuntime { opts.cancellable, std::move(opts.name)); } + // Spawn a coroutine on the SAME strand as the currently-running coroutine + // (identified by its Context). Use this when a child coroutine shares + // mutable state with its parent and must be serialized with it: because both + // run on the one strand, Asio never runs them concurrently, even in + // multi-threaded mode, so the shared state needs no mutex. They still + // interleave cooperatively at every async suspension point, so a child that + // must make progress while the parent is parked in a wait (e.g. a retirement + // reader draining backpressure while its transmitter is blocked on POLLOUT) + // is not starved. Like SpawnOnNewStrand the executor is a strand, so the + // WaitFdReady/WaitReadable handshake requirement is satisfied. + void SpawnOnCurrentStrand(Context ctx, std::function fn, + SpawnOptions opts = {}) { + SpawnTracked(ctx.get_executor(), std::move(fn), opts.cancellable, + std::move(opts.name)); + } + // Run the io_context on `num_threads` threads (the calling thread plus // num_threads-1 helpers). Blocks until the io_context runs out of work or is // stopped, exactly like the single-threaded ioc_.run(). @@ -314,6 +330,13 @@ class AsyncRuntime { Spawn(std::move(fn), std::move(opts)); } + // The co backend is single-threaded, so there is no strand to share: this is + // just a normal spawn. + void SpawnOnCurrentStrand(Context /*ctx*/, std::function fn, + SpawnOptions opts = {}) { + Spawn(std::move(fn), std::move(opts)); + } + // The co scheduler is single-threaded; num_threads is ignored. void Run(int /*num_threads*/ = 1) { scheduler_.Run(); } void Stop() { scheduler_.Stop(); } diff --git a/common/async/socket.cc b/common/async/socket.cc index d225db0a..ac23a04e 100644 --- a/common/async/socket.cc +++ b/common/async/socket.cc @@ -119,10 +119,12 @@ absl::Status WaitFdReady(Context yield, int fd, // to the coroutine's cancellation slot and wait on the notifier with an // *unconnected* slot so asio installs no per-operation timer cancellation. auto exec = yield.get_executor(); - yield.get_cancellation_slot().assign( - [exec, wake](boost::asio::cancellation_type) { - boost::asio::post(exec, [wake]() { wake(); }); - }); + auto cancel_slot = yield.get_cancellation_slot(); + if (cancel_slot.is_connected()) { + cancel_slot.assign([exec, wake](boost::asio::cancellation_type) { + boost::asio::post(exec, [wake]() { wake(); }); + }); + } boost::system::error_code ec; notifier->async_wait( @@ -132,7 +134,9 @@ absl::Status WaitFdReady(Context yield, int fd, // Clear our handler so a late re-emit cannot poke a stale wake, then cancel // the pending descriptor wait and release the real fd so destroying the // shared descriptor does not close it. - yield.get_cancellation_slot().clear(); + if (cancel_slot.is_connected()) { + cancel_slot.clear(); + } boost::system::error_code ignored; sd->cancel(ignored); sd->release(); diff --git a/common/async/wait.cc b/common/async/wait.cc index d453cf27..d3c708e1 100644 --- a/common/async/wait.cc +++ b/common/async/wait.cc @@ -122,13 +122,15 @@ absl::Status WaitReadable(Context ctx, int fd, std::chrono::nanoseconds timeout) // handler to the coroutine's cancellation slot and wait on the notifier with // an *unconnected* slot so asio installs no per-operation timer cancellation. auto exec = ctx.get_executor(); - ctx.get_cancellation_slot().assign( - [exec, wake, cancelled](boost::asio::cancellation_type) { - boost::asio::post(exec, [wake, cancelled]() { - *cancelled = true; - wake(); - }); + auto cancel_slot = ctx.get_cancellation_slot(); + if (cancel_slot.is_connected()) { + cancel_slot.assign([exec, wake, cancelled](boost::asio::cancellation_type) { + boost::asio::post(exec, [wake, cancelled]() { + *cancelled = true; + wake(); }); + }); + } boost::system::error_code ec; notifier->async_wait( @@ -139,7 +141,9 @@ absl::Status WaitReadable(Context ctx, int fd, std::chrono::nanoseconds timeout) // cancellation handler so a late re-emit cannot poke a stale wake, then cancel // any still-pending descriptor/timer waits and release the real fd so // destroying the shared descriptor does not close it. - ctx.get_cancellation_slot().clear(); + if (cancel_slot.is_connected()) { + cancel_slot.clear(); + } boost::system::error_code ignored; sd->cancel(ignored); sd->release(); @@ -158,8 +162,23 @@ absl::Status WaitReadable(Context ctx, int fd, std::chrono::nanoseconds timeout) return absl::InternalError("WaitReadable: operation canceled"); } -absl::StatusOr WaitEither(Context ctx, int fd1, int fd2) { +absl::StatusOr WaitEither(Context ctx, int fd1, int fd2, + std::chrono::nanoseconds timeout) { boost::asio::io_context &ioc = IocFromYield(ctx); + // Fast path: if either fd is already readable, return immediately without + // suspending (see WaitReadable for why real-fd waits need this). + { + struct pollfd fds[2] = {{.fd = fd1, .events = POLLIN, .revents = 0}, + {.fd = fd2, .events = POLLIN, .revents = 0}}; + if (::poll(fds, 2, 0) > 0) { + if ((fds[0].revents & (POLLIN | POLLHUP)) != 0) { + return fd1; + } + if ((fds[1].revents & (POLLIN | POLLHUP)) != 0) { + return fd2; + } + } + } // Register the caller's real fds (no ::dup(); see WaitReadable) and release() // them before the descriptors are destroyed so we never close them. auto sd1 = std::make_shared(ioc, fd1); @@ -167,6 +186,8 @@ absl::StatusOr WaitEither(Context ctx, int fd1, int fd2) { auto result_fd = std::make_shared(-1); auto done = std::make_shared(false); + auto timed_out = std::make_shared(false); + auto cancelled = std::make_shared(false); auto notified = std::make_shared(false); // Shared so the cancellation handler can keep it alive past this frame. auto notifier = std::make_shared(ioc); @@ -182,6 +203,11 @@ absl::StatusOr WaitEither(Context ctx, int fd1, int fd2) { } }; + std::shared_ptr timer; + if (timeout.count() != 0) { + timer = std::make_shared(ioc, timeout); + } + // Bind both descriptor completions to the coroutine's executor (its strand) // so they are serialized with the notifier->async_wait() registration below // and with each other; otherwise a multi-threaded io_context could run a @@ -190,13 +216,17 @@ absl::StatusOr WaitEither(Context ctx, int fd1, int fd2) { sd1->async_wait(boost::asio::posix::stream_descriptor::wait_read, boost::asio::bind_executor( ctx.get_executor(), - [sd1, sd2, result_fd, done, wake, + [sd1, sd2, timer, result_fd, done, timed_out, cancelled, + wake, fd1](const boost::system::error_code &ec) { - if (!*done && !ec) { + if (!*done && !*timed_out && !*cancelled && !ec) { *done = true; *result_fd = fd1; boost::system::error_code ignored; sd2->cancel(ignored); + if (timer != nullptr) { + timer->cancel(); + } } wake(); })); @@ -204,24 +234,47 @@ absl::StatusOr WaitEither(Context ctx, int fd1, int fd2) { sd2->async_wait(boost::asio::posix::stream_descriptor::wait_read, boost::asio::bind_executor( ctx.get_executor(), - [sd1, sd2, result_fd, done, wake, + [sd1, sd2, timer, result_fd, done, timed_out, cancelled, + wake, fd2](const boost::system::error_code &ec) { - if (!*done && !ec) { + if (!*done && !*timed_out && !*cancelled && !ec) { *done = true; *result_fd = fd2; boost::system::error_code ignored; sd1->cancel(ignored); + if (timer != nullptr) { + timer->cancel(); + } } wake(); })); + if (timer != nullptr) { + timer->async_wait(boost::asio::bind_executor( + ctx.get_executor(), [sd1, sd2, done, timed_out, cancelled, + wake](const boost::system::error_code &ec) { + if (!*done && !*cancelled && !ec) { + *timed_out = true; + boost::system::error_code ignored; + sd1->cancel(ignored); + sd2->cancel(ignored); + } + wake(); + })); + } + // Graceful shutdown routes through the same idempotent wake() rather than // cancelling the notifier op directly (see WaitReadable). auto exec = ctx.get_executor(); - ctx.get_cancellation_slot().assign( - [exec, wake](boost::asio::cancellation_type) { - boost::asio::post(exec, [wake]() { wake(); }); + auto cancel_slot = ctx.get_cancellation_slot(); + if (cancel_slot.is_connected()) { + cancel_slot.assign([exec, wake, cancelled](boost::asio::cancellation_type) { + boost::asio::post(exec, [wake, cancelled]() { + *cancelled = true; + wake(); }); + }); + } boost::system::error_code ec; notifier->async_wait( @@ -230,14 +283,22 @@ absl::StatusOr WaitEither(Context ctx, int fd1, int fd2) { // Clear our handler so a late re-emit cannot poke a stale wake, then release // the real fds so destroying the shared descriptors does not close them. - ctx.get_cancellation_slot().clear(); + if (cancel_slot.is_connected()) { + cancel_slot.clear(); + } boost::system::error_code ignored; sd1->cancel(ignored); sd2->cancel(ignored); sd1->release(); sd2->release(); + if (timer != nullptr) { + timer->cancel(); + } if (*result_fd < 0) { + if (*timed_out) { + return absl::DeadlineExceededError("Timeout waiting for fd"); + } return absl::InternalError("WaitEither: no fd became ready"); } return *result_fd; @@ -259,15 +320,19 @@ void Sleep(Context ctx, std::chrono::nanoseconds duration) { } }; auto exec = ctx.get_executor(); - ctx.get_cancellation_slot().assign( - [exec, wake](boost::asio::cancellation_type) { - boost::asio::post(exec, [wake]() { wake(); }); - }); + auto cancel_slot = ctx.get_cancellation_slot(); + if (cancel_slot.is_connected()) { + cancel_slot.assign([exec, wake](boost::asio::cancellation_type) { + boost::asio::post(exec, [wake]() { wake(); }); + }); + } boost::system::error_code ec; timer->async_wait( boost::asio::bind_cancellation_slot(boost::asio::cancellation_slot(), ctx[ec])); - ctx.get_cancellation_slot().clear(); + if (cancel_slot.is_connected()) { + cancel_slot.clear(); + } } } // namespace subspace::async @@ -284,9 +349,14 @@ absl::Status WaitReadable(Context ctx, int fd, std::chrono::nanoseconds timeout) return absl::OkStatus(); } -absl::StatusOr WaitEither(Context ctx, int fd1, int fd2) { - int r = ctx->Wait(std::vector{fd1, fd2}, POLLIN); +absl::StatusOr WaitEither(Context ctx, int fd1, int fd2, + std::chrono::nanoseconds timeout) { + int r = ctx->Wait(std::vector{fd1, fd2}, POLLIN, + static_cast(timeout.count())); if (r == -1) { + if (timeout.count() != 0) { + return absl::DeadlineExceededError("Timeout waiting for fd"); + } return absl::InternalError("WaitEither: no fd became ready"); } return r; diff --git a/common/async/wait.h b/common/async/wait.h index 5075b76d..16cd0b4c 100644 --- a/common/async/wait.h +++ b/common/async/wait.h @@ -25,8 +25,11 @@ absl::Status WaitReadable(Context ctx, int fd, std::chrono::nanoseconds(0)); // Wait until either `fd1` or `fd2` becomes readable. Returns the raw fd that -// became ready. -absl::StatusOr WaitEither(Context ctx, int fd1, int fd2); +// became ready. A non-zero timeout returns DeadlineExceededError if neither fd +// is ready in time; a zero timeout waits forever. +absl::StatusOr WaitEither(Context ctx, int fd1, int fd2, + std::chrono::nanoseconds timeout = + std::chrono::nanoseconds(0)); // Suspend the current coroutine for `duration`. void Sleep(Context ctx, std::chrono::nanoseconds duration); diff --git a/coro_rpc/server/rpc_server.cc b/coro_rpc/server/rpc_server.cc index ca7694fd..fcda6d9b 100644 --- a/coro_rpc/server/rpc_server.cc +++ b/coro_rpc/server/rpc_server.cc @@ -8,6 +8,7 @@ #include "proto/subspace.pb.h" #include #include +#include namespace subspace::coro_rpc { @@ -751,27 +752,36 @@ boost::asio::awaitable RpcServer::SessionStreamingMethodCoroutine( method_instance->method->name.c_str()); AnyStreamWriter writer(server, session, method_instance, request); + auto stop_pipe_or = toolbelt::Pipe::Create(); + if (!stop_pipe_or.ok()) { + server->logger_.Log(toolbelt::LogLevel::kError, + "Failed to create stream cancellation stop pipe: %s", + stop_pipe_or.status().ToString().c_str()); + continue; + } + auto stop_pipe = + std::make_shared(std::move(*stop_pipe_or)); + auto writer_state = writer.state; + const int32_t session_id = session->session_id; + const int32_t request_id = request.request_id(); // Spawn a coroutine to read the cancellation channel. boost::asio::co_spawn( *server->io_context_, - [server, session, method_instance, &writer, - &request]() -> boost::asio::awaitable { - int dup_interrupt = - ::dup(server->interrupt_pipe_.ReadFd().Fd()); - toolbelt::FileDescriptor interrupt(dup_interrupt); - while (!writer.IsCancelled()) { + [server, method_instance, stop_pipe, writer_state, session_id, + request_id]() -> boost::asio::awaitable { + while (!writer_state->is_cancelled.load(std::memory_order_acquire)) { int cancel_fd = method_instance->cancel_subscriber->GetPollFd().fd; auto s = - co_await async_wait_either(cancel_fd, interrupt.Fd()); + co_await async_wait_either(cancel_fd, stop_pipe->ReadFd().Fd()); if (!s.ok()) { server->logger_.Log(toolbelt::LogLevel::kError, "Error waiting for cancel: %s", s.status().ToString().c_str()); co_return; } - if (*s == interrupt.Fd()) { + if (*s == stop_pipe->ReadFd().Fd()) { break; } bool cancel_ok = false; @@ -793,14 +803,14 @@ boost::asio::awaitable RpcServer::SessionStreamingMethodCoroutine( msg.status().ToString().c_str()); continue; } - if (cancel.session_id() == session->session_id && - cancel.request_id() == request.request_id()) { + if (cancel.session_id() == session_id && + cancel.request_id() == request_id) { cancel_ok = true; break; } } if (cancel_ok) { - writer.Cancel(); + writer_state->is_cancelled.store(true, std::memory_order_release); } } }, @@ -820,6 +830,8 @@ boost::asio::awaitable RpcServer::SessionStreamingMethodCoroutine( method_instance->method->name, method_status.ToString())); } + char stop = 1; + (void)::write(stop_pipe->WriteFd().Fd(), &stop, 1); } } diff --git a/coro_rpc/server/rpc_server.h b/coro_rpc/server/rpc_server.h index 090d0e08..386c2c0c 100644 --- a/coro_rpc/server/rpc_server.h +++ b/coro_rpc/server/rpc_server.h @@ -17,6 +17,7 @@ #include #include #include +#include namespace subspace::coro_rpc { @@ -27,26 +28,33 @@ namespace internal { struct Session; struct MethodInstance; +struct StreamWriterState { + std::atomic_bool is_cancelled{false}; +}; + struct AnyStreamWriter { AnyStreamWriter(std::shared_ptr server, std::shared_ptr session, std::shared_ptr method_instance, const RpcRequest &request) : server(std::move(server)), session(std::move(session)), - method_instance(std::move(method_instance)), request(request) {} + method_instance(std::move(method_instance)), request(request), + state(std::make_shared()) {} boost::asio::awaitable Write(std::unique_ptr res); boost::asio::awaitable Finish(); - void Cancel() { is_cancelled = true; } + void Cancel() { state->is_cancelled.store(true, std::memory_order_release); } - bool IsCancelled() const { return is_cancelled; } + bool IsCancelled() const { + return state->is_cancelled.load(std::memory_order_acquire); + } std::shared_ptr server; std::shared_ptr session; std::shared_ptr method_instance; - const RpcRequest &request; - bool is_cancelled = false; + RpcRequest request; + std::shared_ptr state; }; // An item pushed by either the reply function or the error function of an diff --git a/coro_rpc/server/server_test.cc b/coro_rpc/server/server_test.cc index 243f59e2..a1943921 100644 --- a/coro_rpc/server/server_test.cc +++ b/coro_rpc/server/server_test.cc @@ -154,6 +154,18 @@ static std::shared_ptr BuildServer() { return server; } +TEST_F(CoroServerTest, StreamWriterOwnsRequestMetadata) { + std::unique_ptr writer; + { + subspace::RpcRequest request; + request.set_request_id(1234); + writer = std::make_unique( + nullptr, nullptr, nullptr, request); + } + + EXPECT_EQ(1234, writer->request.request_id()); +} + struct ServerContext { std::shared_ptr client; std::shared_ptr pub; diff --git a/rpc/server/rpc_server.cc b/rpc/server/rpc_server.cc index 29994337..732631bb 100644 --- a/rpc/server/rpc_server.cc +++ b/rpc/server/rpc_server.cc @@ -756,23 +756,35 @@ void RpcServer::SessionStreamingMethodCoroutine( method_instance->method->name.c_str()); AnyStreamWriter writer(server, session, method_instance, request); + auto stop_pipe_or = toolbelt::Pipe::Create(); + if (!stop_pipe_or.ok()) { + server->logger_.Log(toolbelt::LogLevel::kError, + "Failed to create stream cancellation stop pipe: %s", + stop_pipe_or.status().ToString().c_str()); + continue; + } + auto stop_pipe = + std::make_shared(std::move(*stop_pipe_or)); + auto writer_state = writer.state; + const int32_t session_id = session->session_id; + const int32_t request_id = request.request_id(); // Start a coroutine to read the cancellation channel and cancel the // StreamWriter if a cancellation request is received. server->AddCoroutine(std::make_unique( - *server->scheduler_, [server, session, method_instance, &writer, - &request](co::Coroutine *c) { - toolbelt::FileDescriptor interrupt( - dup(server->interrupt_pipe_.ReadFd().Fd())); - while (!writer.IsCancelled()) { - auto s = method_instance->cancel_subscriber->Wait(interrupt, c); + *server->scheduler_, + [server, method_instance, stop_pipe, writer_state, session_id, + request_id](co::Coroutine *c) { + while (!writer_state->is_cancelled.load(std::memory_order_acquire)) { + auto s = + method_instance->cancel_subscriber->Wait(stop_pipe->ReadFd(), c); if (!s.ok()) { server->logger_.Log(toolbelt::LogLevel::kError, "Error waiting for cancel: %s", s.status().ToString().c_str()); return; } - if (*s == interrupt.Fd()) { + if (*s == stop_pipe->ReadFd().Fd()) { break; } bool request_ok = false; @@ -795,14 +807,14 @@ void RpcServer::SessionStreamingMethodCoroutine( msg.status().ToString().c_str()); continue; } - if (cancel.session_id() == session->session_id && - cancel.request_id() == request.request_id()) { + if (cancel.session_id() == session_id && + cancel.request_id() == request_id) { request_ok = true; break; } } if (request_ok) { - writer.Cancel(); + writer_state->is_cancelled.store(true, std::memory_order_release); } } })); @@ -823,6 +835,8 @@ void RpcServer::SessionStreamingMethodCoroutine( method_status.ToString()), c); } + char stop = 1; + (void)::write(stop_pipe->WriteFd().Fd(), &stop, 1); } } diff --git a/rpc/server/rpc_server.h b/rpc/server/rpc_server.h index 0a3cd936..db9155de 100644 --- a/rpc/server/rpc_server.h +++ b/rpc/server/rpc_server.h @@ -12,6 +12,8 @@ #include "toolbelt/logging.h" #include "toolbelt/pipe.h" +#include + namespace subspace { class RpcServer; @@ -21,27 +23,34 @@ namespace internal { struct Session; struct MethodInstance; +struct StreamWriterState { + std::atomic_bool is_cancelled{false}; +}; + struct AnyStreamWriter { AnyStreamWriter(std::shared_ptr server, std::shared_ptr session, std::shared_ptr method_instance, const RpcRequest &request) : server(std::move(server)), session(std::move(session)), - method_instance(std::move(method_instance)), request(request) {} + method_instance(std::move(method_instance)), request(request), + state(std::make_shared()) {} // Returns true if the write worked, false if the request was cancelled. bool Write(std::unique_ptr res, co::Coroutine *c); void Finish(co::Coroutine *c); - void Cancel() { is_cancelled = true; } + void Cancel() { state->is_cancelled.store(true, std::memory_order_release); } - bool IsCancelled() const { return is_cancelled; } + bool IsCancelled() const { + return state->is_cancelled.load(std::memory_order_acquire); + } std::shared_ptr server; std::shared_ptr session; std::shared_ptr method_instance; - const RpcRequest &request; - bool is_cancelled = false; + RpcRequest request; + std::shared_ptr state; }; // An item pushed by either the reply function or the error function of an diff --git a/rpc/server/server_test.cc b/rpc/server/server_test.cc index 578a7044..077132f2 100644 --- a/rpc/server/server_test.cc +++ b/rpc/server/server_test.cc @@ -237,6 +237,18 @@ static std::shared_ptr BuildServer() { return server; } +TEST_F(ServerTest, StreamWriterOwnsRequestMetadata) { + std::unique_ptr writer; + { + subspace::RpcRequest request; + request.set_request_id(1234); + writer = std::make_unique( + nullptr, nullptr, nullptr, request); + } + + EXPECT_EQ(1234, writer->request.request_id()); +} + struct ServerContext { std::shared_ptr client; std::shared_ptr pub; diff --git a/server/server.cc b/server/server.cc index 8cd07e63..f3ea5b6d 100644 --- a/server/server.cc +++ b/server/server.cc @@ -49,6 +49,57 @@ static bool ChannelUsesSplitBuffers(const ServerChannel *channel) { split_channel->GetSplitBufferOptions().use_split_buffers; } +// Tracks the ActiveMessages that a bridge transmitter has sent but whose slots +// the peer has not yet retired, indexed by slot_id. It is touched by two +// coroutines - the transmitter (Track/Release) and its retirement reader +// (Release) - but both are spawned on the same strand (see +// SpawnOnCurrentStrand at the retirement-reader spawn site), so their accesses +// are serialized by the strand and need no mutex, even when the io_context runs +// on multiple threads. +struct Server::BridgeRetirementState { + explicit BridgeRetirementState(size_t num_slots) + : active_messages(num_slots) {} + + ~BridgeRetirementState() { + for (auto &active : active_messages) { + if (active != nullptr) { + active->DecRef(); + } + } + } + + void Track(int slot_id, std::shared_ptr active) { + if (active == nullptr) { + return; + } + if (slot_id < 0 || + static_cast(slot_id) >= active_messages.size()) { + active->DecRef(); + return; + } + std::shared_ptr previous = + std::move(active_messages[slot_id]); + active_messages[slot_id] = std::move(active); + if (previous != nullptr) { + previous->DecRef(); + } + } + + void Release(int slot_id) { + if (slot_id < 0 || + static_cast(slot_id) >= active_messages.size()) { + return; + } + std::shared_ptr active = + std::move(active_messages[slot_id]); + if (active != nullptr) { + active->DecRef(); + } + } + + std::vector> active_messages; +}; + #if SUBSPACE_SHMEM_MODE == SUBSPACE_SHMEM_MODE_POSIX static bool CleanupSplitBufferMetadataFile(const std::filesystem::path &path) { absl::StatusOr metadata = @@ -386,6 +437,34 @@ BridgeAddressWithPort(const toolbelt::SocketAddress &addr, int port) { } } +absl::StatusOr +ParseChannelAddress(const ChannelAddress &address, const char *field_name) { + constexpr size_t kAddressSize = sizeof(uint32_t); + if (address.address().size() != kAddressSize) { + return absl::InvalidArgumentError(absl::StrFormat( + "%s has invalid address length %zu; expected %zu", field_name, + address.address().size(), kAddressSize)); + } + + switch (address.family()) { + case ChannelAddress::FAMILY_INET: { + in_addr ip_addr; + memcpy(&ip_addr, address.address().data(), sizeof(ip_addr)); + // ChannelAddress stores IPv4 bytes in host order on the wire. + ip_addr.s_addr = ntohl(ip_addr.s_addr); + return toolbelt::SocketAddress(ip_addr, address.port()); + } + case ChannelAddress::FAMILY_VSOCK: { + uint32_t cid; + memcpy(&cid, address.address().data(), sizeof(cid)); + return toolbelt::SocketAddress(cid, address.port()); + } + default: + return absl::InvalidArgumentError(absl::StrFormat( + "%s has unsupported address family %d", field_name, address.family())); + } +} + #if SUBSPACE_CORO_BACKEND == SUBSPACE_CORO_BACKEND_CO Server::Server(co::CoroutineScheduler &scheduler, const std::string &socket_name, const std::string &interface, @@ -1309,6 +1388,7 @@ absl::Status Server::RecoverFromShadow(RecoveredState &state) { channel->AddUser(rsub.id, std::move(sub)); } + channel->RegisterExistingSubscribers(); channels_.emplace(rch.name, channel); logger_.Log(toolbelt::LogLevel::kInfo, @@ -1859,7 +1939,15 @@ void Server::BridgeTransmitterCoroutine(async::Context ctx, subscribed.set_split_buffers(wire_split_buffers); subscribed.set_split_buffers_over_bridge(info.split_buffers_over_bridge); - async::StreamSocket retirement_listener; + std::shared_ptr retirement_listener; + struct CloseRetirementListenerOnExit { + std::shared_ptr &socket; + ~CloseRetirementListenerOnExit() { + if (socket != nullptr) { + socket->Close(); + } + } + } close_retirement_listener_on_exit{retirement_listener}; toolbelt::SocketAddress retirement_addr; if (notify_retirement) { @@ -1867,15 +1955,16 @@ void Server::BridgeTransmitterCoroutine(async::Context ctx, subscribed.set_notify_retirement(true); // Allocate a listen socket to wait for an incoming connection from the // other server. + retirement_listener = std::make_shared(); - absl::Status s = BindBridgeListener(retirement_listener); + absl::Status s = BindBridgeListener(*retirement_listener); if (!s.ok()) { logger_.Log(toolbelt::LogLevel::kError, "Unable to bind socket for retirement receiver for %s: %s", channel_name.c_str(), s.ToString().c_str()); return; } - retirement_addr = retirement_listener.BoundAddress(); + retirement_addr = retirement_listener->BoundAddress(); logger_.Log(toolbelt::LogLevel::kDebug, "Retirement listener: %s", retirement_addr.ToString().c_str()); @@ -1944,19 +2033,24 @@ void Server::BridgeTransmitterCoroutine(async::Context ctx, // send over the bridge. This extends the lifetime of the AciveMessage until // the other side of the bridge sends us a retirement notification for the // message's slot. - std::shared_ptr>> - active_retirement_msgs = - std::make_shared>>(); + std::shared_ptr retirement_state; bool notifying_of_retirement = notify_retirement; - // Spawn a coroutine to read from the retirement connection. + // Spawn a coroutine to read from the retirement connection. It runs on this + // transmitter's own strand (not a fresh one) so that its Release calls are + // serialized with our Track/Release on the shared retirement_state without a + // mutex. The two still interleave cooperatively: whenever this transmitter + // is parked in a wait (subscriber trigger, or POLLOUT under bridge + // backpressure) the strand is free to run the retirement reader, so incoming + // retirements keep relieving that backpressure. if (notifying_of_retirement) { - active_retirement_msgs->resize(info.num_slots); - runtime_.SpawnOnNewStrand( - [this, &retirement_listener, - active_retirement_msgs](async::Context ret_ctx) mutable { + retirement_state = std::make_shared(info.num_slots); + runtime_.SpawnOnCurrentStrand( + ctx, + [this, retirement_listener, + retirement_state](async::Context ret_ctx) mutable { return RetirementReceiverCoroutine(ret_ctx, retirement_listener, - active_retirement_msgs); + retirement_state); }, {.name = absl::StrFormat("Retirement listener for %s", channel_name), .interrupt_fd = shutdown_trigger_fd_.GetPollFd().Fd()}); @@ -2037,22 +2131,26 @@ void Server::BridgeTransmitterCoroutine(async::Context ctx, // The backpressure received here will be applied upwards because // we will stop reading the messages from the channel and thus // backpressure any publishers writing to that channel. + const int32_t slot_id = msg->slot_id; + if (notifying_of_retirement) { + // Track before sending: the peer can consume and retire the bridged + // slot on the separate retirement socket as soon as SendBridgeMessage + // returns. + retirement_state->Track(slot_id, std::move(msg->active_message)); + } if (absl::Status status = SendBridgeMessage(ctx, bridge, *msg, prefix, prefix_area, wire_split_buffers); !status.ok()) { + if (notifying_of_retirement) { + retirement_state->Release(slot_id); + } done = true; logger_.Log(toolbelt::LogLevel::kError, "Failed to send bridge message for %s: %s", channel_name.c_str(), status.ToString().c_str()); break; } - if (notifying_of_retirement) { - // We need to keep track of the message so that we can retire it - // when the other side retires the slot. - (*active_retirement_msgs)[msg->slot_id] = - std::move(msg->active_message); - } } const ChannelCounters &counters = sub->GetCounters(); @@ -2073,12 +2171,12 @@ void Server::BridgeTransmitterCoroutine(async::Context ctx, // If these references are the last reference to the slot, it will be retired // on this side and the publisher's retirement FD is sent the slot. void Server::RetirementReceiverCoroutine( - async::Context ctx, async::StreamSocket &retirement_listener, - std::shared_ptr>> - active_retirement_msgs) { + async::Context ctx, + std::shared_ptr retirement_listener, + std::shared_ptr retirement_state) { // Accept connection on the retirement listener socket. absl::StatusOr retirement_socket = - retirement_listener.Accept(ctx); + retirement_listener->Accept(ctx); if (!retirement_socket.ok()) { logger_.Log(toolbelt::LogLevel::kError, "Failed to accept retirement connection: %s", @@ -2109,11 +2207,7 @@ void Server::RetirementReceiverCoroutine( // being sent (which is serialized as 0 bytes. Remove the adustment. slot_id -= 1; - if (slot_id < 0 || - static_cast(slot_id) >= active_retirement_msgs->size()) { - continue; - } - (*active_retirement_msgs)[slot_id]->DecRef(); + retirement_state->Release(slot_id); } } @@ -2281,43 +2375,30 @@ void Server::BridgeReceiverCoroutine(async::Context ctx, return; } - std::unique_ptr retirement_transmitter; + std::shared_ptr retirement_transmitter; + struct CloseRetirementTransmitterOnExit { + std::shared_ptr &socket; + ~CloseRetirementTransmitterOnExit() { + if (socket != nullptr) { + socket->Close(); + } + } + } close_retirement_transmitter_on_exit{retirement_transmitter}; if (subscribed.notify_retirement()) { - retirement_transmitter = std::make_unique(); - toolbelt::SocketAddress retirement_addr; + retirement_transmitter = std::make_shared(); // The address family travels in the message (see SendSubscribeMessage); // unset defaults to FAMILY_INET for compatibility with older peers. - switch (subscribed.retirement_socket().family()) { - case ChannelAddress::FAMILY_INET: { - // IPv4 address. - struct sockaddr_in tmp_addr; - memset(&tmp_addr, 0, sizeof(tmp_addr)); - tmp_addr.sin_family = AF_INET; - tmp_addr.sin_port = htons(subscribed.retirement_socket().port()); - tmp_addr.sin_addr.s_addr = ntohl(*reinterpret_cast( - subscribed.retirement_socket().address().data())); - retirement_addr = toolbelt::SocketAddress(tmp_addr); - break; - } - case ChannelAddress::FAMILY_VSOCK: { - // Virtual address. - struct sockaddr_vm tmp_addr; - memset(&tmp_addr, 0, sizeof(tmp_addr)); - tmp_addr.svm_family = AF_VSOCK; - tmp_addr.svm_port = subscribed.retirement_socket().port(); - tmp_addr.svm_cid = *reinterpret_cast( - subscribed.retirement_socket().address().data()); - retirement_addr = toolbelt::SocketAddress(tmp_addr); - break; - } - default: + absl::StatusOr retirement_addr = + ParseChannelAddress(subscribed.retirement_socket(), "retirement socket"); + if (!retirement_addr.ok()) { logger_.Log(toolbelt::LogLevel::kError, - "Unsupported address family for retirement notification"); + "Invalid retirement socket for %s: %s", channel_name.c_str(), + retirement_addr.status().ToString().c_str()); return; } - s = retirement_transmitter->Connect(retirement_addr); + s = retirement_transmitter->Connect(*retirement_addr); if (!s.ok()) { logger_.Log(toolbelt::LogLevel::kError, "Failed to connect to retirement socket for %s: %s", @@ -2366,9 +2447,10 @@ void Server::BridgeReceiverCoroutine(async::Context ctx, // the socket connected to the other server. runtime_.SpawnOnNewStrand( [this, retirement_fd = std::move(retirement_fd), - &retirement_transmitter, channel_name](async::Context ret_ctx) mutable { + retirement_transmitter, + channel_name](async::Context ret_ctx) mutable { RetirementCoroutine(ret_ctx, channel_name, std::move(retirement_fd), - std::move(retirement_transmitter)); + retirement_transmitter); }, {.name = absl::StrFormat("Retirement notifier for %s", channel_name), .interrupt_fd = shutdown_trigger_fd_.GetPollFd().Fd()}); @@ -2471,7 +2553,7 @@ void Server::BridgeReceiverCoroutine(async::Context ctx, void Server::RetirementCoroutine( async::Context ctx, const std::string &channel_name, toolbelt::FileDescriptor &&retirement_fd, - std::unique_ptr retirement_transmitter) { + std::shared_ptr retirement_transmitter) { logger_.Log(toolbelt::LogLevel::kDebug, "Retirement coroutine for %s running", channel_name.c_str()); #if SUBSPACE_CORO_BACKEND == SUBSPACE_CORO_BACKEND_ASIO @@ -2650,28 +2732,13 @@ void Server::IncomingSubscribe(const Discovery::Subscribe &subscribe, // carried in the message itself (FAMILY_INET for TCP, FAMILY_VSOCK for // vsock) so the two servers do not have to share the same local address // type. Older peers leave family unset, which defaults to FAMILY_INET. - toolbelt::SocketAddress subscriber_addr; - switch (subscribe.receiver().family()) { - case ChannelAddress::FAMILY_INET: { - in_addr subscriber_ip; - memcpy(&subscriber_ip, subscribe.receiver().address().data(), - sizeof(subscriber_ip)); - // Need this in host byte order. - subscriber_ip.s_addr = ntohl(subscriber_ip.s_addr); - subscriber_addr = - toolbelt::SocketAddress(subscriber_ip, subscribe.receiver().port()); - break; - } - case ChannelAddress::FAMILY_VSOCK: { - uint32_t cid = *reinterpret_cast( - subscribe.receiver().address().data()); - subscriber_addr = - toolbelt::SocketAddress(cid, subscribe.receiver().port()); - break; - } - default: + absl::StatusOr subscriber_addr = + ParseChannelAddress(subscribe.receiver(), "subscribe receiver"); + if (!subscriber_addr.ok()) { logger_.Log(toolbelt::LogLevel::kError, - "Unknown address family for subscribe receiver"); + "Invalid subscribe receiver for %s: %s", + subscribe.channel_name().c_str(), + subscriber_addr.status().ToString().c_str()); return; } @@ -2690,7 +2757,7 @@ void Server::IncomingSubscribe(const Discovery::Subscribe &subscribe, }; runtime_.SpawnOnNewStrand( [this, info = std::move(info), pub_reliable, sub_reliable, - subscriber_addr = std::move(subscriber_addr), sender, + subscriber_addr = std::move(*subscriber_addr), sender, notify_retirement](async::Context ctx) mutable { BridgeTransmitterCoroutine(ctx, std::move(info), pub_reliable, sub_reliable, std::move(subscriber_addr), diff --git a/server/server.h b/server/server.h index 59e8efd2..6341fbb6 100644 --- a/server/server.h +++ b/server/server.h @@ -44,6 +44,9 @@ constexpr int64_t kServerWaiting = 3; void ClosePluginsOnShutdown(); bool ShouldClosePluginsOnShutdown(); +absl::StatusOr +ParseChannelAddress(const ChannelAddress &address, const char *field_name); + // The Subspace server. // This is a single-threaded, coroutine-based server that maintains shared // memory IPC channels and communicates with other servers to allow for @@ -300,6 +303,7 @@ class Server { bool wire_split_buffers = false; bool split_buffers_over_bridge = false; }; + struct BridgeRetirementState; void BridgeTransmitterCoroutine(async::Context ctx, BridgeChannelInfo info, bool pub_reliable, bool sub_reliable, toolbelt::SocketAddress subscriber, @@ -314,11 +318,12 @@ class Server { void RetirementCoroutine( async::Context ctx, const std::string &channel_name, toolbelt::FileDescriptor &&retirement_fd, - std::unique_ptr retirement_transmitter); + std::shared_ptr retirement_transmitter); void RetirementReceiverCoroutine( - async::Context ctx, async::StreamSocket &retirement_listener, - std::shared_ptr>> active_retirement_msgs); + async::Context ctx, + std::shared_ptr retirement_listener, + std::shared_ptr retirement_state); void SubscribeOverBridge(ServerChannel *channel, bool reliable, bool split_buffers_over_bridge, diff --git a/server/server_test.cc b/server/server_test.cc index 15f8f269..32ca3303 100644 --- a/server/server_test.cc +++ b/server/server_test.cc @@ -138,6 +138,47 @@ TEST_F(ServerTest, InitSuccess) { ASSERT_FALSE(fds.empty()); } +TEST(ChannelAddressParsingTest, ParsesValidInetAddress) { + subspace::ChannelAddress address; + in_addr ip_addr; + ip_addr.s_addr = htonl(INADDR_LOOPBACK); + address.set_address(&ip_addr, sizeof(ip_addr)); + address.set_port(12345); + address.set_family(subspace::ChannelAddress::FAMILY_INET); + + absl::StatusOr parsed = + subspace::ParseChannelAddress(address, "test address"); + ASSERT_OK(parsed); + EXPECT_EQ(parsed->Type(), toolbelt::SocketAddress::kAddressInet); + EXPECT_EQ(parsed->Port(), 12345); +} + +TEST(ChannelAddressParsingTest, RejectsShortInetAddress) { + subspace::ChannelAddress address; + address.set_address("x"); + address.set_port(12345); + address.set_family(subspace::ChannelAddress::FAMILY_INET); + + absl::StatusOr parsed = + subspace::ParseChannelAddress(address, "test address"); + ASSERT_FALSE(parsed.ok()); + EXPECT_THAT(parsed.status().message(), + ::testing::HasSubstr("invalid address length")); +} + +TEST(ChannelAddressParsingTest, RejectsShortVsockAddress) { + subspace::ChannelAddress address; + address.set_address("x"); + address.set_port(12345); + address.set_family(subspace::ChannelAddress::FAMILY_VSOCK); + + absl::StatusOr parsed = + subspace::ParseChannelAddress(address, "test address"); + ASSERT_FALSE(parsed.ok()); + EXPECT_THAT(parsed.status().message(), + ::testing::HasSubstr("invalid address length")); +} + TEST_F(ServerTest, CreatePublisherSuccess) { RawConnection conn; ASSERT_OK(conn.Connect(Socket())); diff --git a/shadow/shadow_test.cc b/shadow/shadow_test.cc index b05eef85..7b7ef9d5 100644 --- a/shadow/shadow_test.cc +++ b/shadow/shadow_test.cc @@ -561,6 +561,64 @@ TEST_F(ShadowRecoveryTest, ServerRecoversStateFromShadow) { StopShadow(); } +TEST_F(ShadowRecoveryTest, RecoveryRebuildsCcbSubscriberRegistrations) { + signal(SIGPIPE, SIG_IGN); + + StartShadow(); + StartServer(); + + constexpr char kChannel[] = "recovery_subscriber_registrations"; + + subspace::Client client; + client.SetThreadSafe(true); + ASSERT_THAT(client.Init(RecoveryServerSocket()), IsOk()); + + auto pub = client.CreatePublisher(kChannel, 256, 8); + ASSERT_THAT(pub, IsOk()); + + auto sub = client.CreateSubscriber(kChannel); + ASSERT_THAT(sub, IsOk()); + + ASSERT_TRUE(WaitForShadowState([this]() { + return shadow_->WithChannels([](auto &channels) { + auto it = channels.find("recovery_subscriber_registrations"); + return it != channels.end() && it->second.publishers.size() == 1 && + it->second.subscribers.size() == 1; + }); + })); + + subspace::ServerChannel *channel = server_->FindChannel(kChannel); + ASSERT_NE(nullptr, channel); + channel->GetCcb()->subscribers.ClearAll(); + channel->GetCcb()->num_subs = subspace::SubscriberCounter(); + ASSERT_EQ(0, channel->NumSubscribers(-1)); + + server_->ForEachShadow( + [](const std::unique_ptr &s) { s->Close(); }); + StopServer(); + + StartServer(); + + subspace::ServerChannel *recovered = server_->FindChannel(kChannel); + ASSERT_NE(nullptr, recovered); + EXPECT_EQ(1, recovered->NumSubscribers(-1)); + + constexpr char kPayload[] = "after_recovery"; + absl::StatusOr buffer = pub->GetMessageBuffer(); + ASSERT_THAT(buffer, IsOk()); + memcpy(*buffer, kPayload, sizeof(kPayload) - 1); + ASSERT_THAT(pub->PublishMessage(sizeof(kPayload) - 1), IsOk()); + + absl::StatusOr msg = + sub->ReadMessage(subspace::ReadMode::kReadNewest); + ASSERT_THAT(msg, IsOk()); + ASSERT_EQ(msg->length, sizeof(kPayload) - 1); + EXPECT_EQ(memcmp(msg->buffer, kPayload, sizeof(kPayload) - 1), 0); + + StopServer(); + StopShadow(); +} + TEST_F(ShadowRecoveryTest, ServerRecoversSplitBufferStateFromShadow) { signal(SIGPIPE, SIG_IGN);