diff --git a/BUILD.bazel b/BUILD.bazel index 6acb6a77353..9425f146279 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -1,7 +1,7 @@ load("@aspect_rules_js//js:defs.bzl", "js_library") load("@aspect_rules_js//npm:defs.bzl", "npm_link_package") load("@bazel_skylib//lib:selects.bzl", "selects") -load("@bazel_skylib//rules:common_settings.bzl", "bool_flag") +load("@bazel_skylib//rules:common_settings.bzl", "bool_flag", "string_flag") load("@npm//:capnp-es/package_json.bzl", capnp_es_bins = "bin") load("@npm//:defs.bzl", "npm_link_all_packages") load("//:build/wd_cc_embed.bzl", "wd_cc_embed") @@ -128,6 +128,44 @@ bool_flag( build_setting_default = True, ) +# ============================================================================= +# I/O backend selection: a build/compile/link-time hermeticity guarantee. +# +# --//:io_backend=rust (default) -- the rust I/O layer: the process event loop is tokio +# (kj-rs-tokio) and every socket/stream is tokio-backed (kj-rs-io); the +# C++ layers above the streams (kj-http, kj-tls, capnp-rpc) are unchanged +# and run over those tokio streams. In this config the forbidden concrete +# C++ OS-I/O target must be ABSENT from the link, enforced by the +# build-graph aspect in //build:rust_io_backend.bzl (whose forbidden +# set grows as later migration stages move more of the stack to rust). +# --//:io_backend=cxx -- workerd's I/O is the concrete C++ stack end to end, including the +# kj OS event loop and sockets. Byte-identical to the pre-migration +# build. +# +# Code that must differ per backend keys off the WORKERD_RUST_IO_BACKEND define (see +# //src/workerd/util:setup-async-io) or off the config_settings below in a select(). +string_flag( + name = "io_backend", + build_setting_default = "rust", + values = [ + "cxx", + "rust", + ], + visibility = ["//visibility:public"], +) + +config_setting( + name = "io_backend_rust", + flag_values = {":io_backend": "rust"}, + visibility = ["//visibility:public"], +) + +config_setting( + name = "io_backend_cxx", + flag_values = {":io_backend": "cxx"}, + visibility = ["//visibility:public"], +) + config_setting( name = "set_dead_strip", flag_values = {"dead_strip": "True"}, diff --git a/build/deps/deps.jsonc b/build/deps/deps.jsonc index 0445e951f54..59ff6ae93f4 100644 --- a/build/deps/deps.jsonc +++ b/build/deps/deps.jsonc @@ -36,7 +36,9 @@ "type": "github_tarball", "owner": "capnproto", "repo": "capnproto", - "branch": "v2", + // Interim: v2 plus kj::Rc::disown()/reown() (mirror of kj::Arc's), needed by the Rust + // I/O bridge's waker. Switch back to "v2" once the upstream PR merges. + "branch": "dlapid/RcdisReOwn", "extra_strip_prefix": "/c++" }, // We want to avoid version skew with v8, so we use identical versions. Keep this diff --git a/build/deps/gen/deps.MODULE.bazel b/build/deps/gen/deps.MODULE.bazel index b4965010eeb..ec83c56465a 100644 --- a/build/deps/gen/deps.MODULE.bazel +++ b/build/deps/gen/deps.MODULE.bazel @@ -27,10 +27,10 @@ bazel_dep(name = "brotli", version = "1.2.0.bcr.1") # capnp-cpp http.archive( name = "capnp-cpp", - sha256 = "6753378bd099029cb2830fecd32dd158218019e459ffd3c8e379cbf025906eb8", - strip_prefix = "capnproto-capnproto-a1cd1c4/c++", + sha256 = "afaf9a84342533e12c8f8c67adcfeb59c74ba86f04c15499736de453c3edea22", + strip_prefix = "capnproto-capnproto-1ac3c72/c++", type = "tgz", - url = "https://github.com/capnproto/capnproto/tarball/a1cd1c4b3d241b77478035a6ccad8b0fb587d444", + url = "https://github.com/capnproto/capnproto/tarball/1ac3c7259ddf3917d7c5947978e1f97bf328b699", ) use_repo(http, "capnp-cpp") diff --git a/build/kj_test.bzl b/build/kj_test.bzl index 6d4ad136cbd..9adefbac5e4 100644 --- a/build/kj_test.bzl +++ b/build/kj_test.bzl @@ -1,5 +1,6 @@ load("@rules_cc//cc:cc_binary.bzl", "cc_binary") load("@rules_shell//shell:sh_test.bzl", "sh_test") +load("//:build/rust_io_backend.bzl", "rust_io_backend_local_defines") def kj_test( src, @@ -7,6 +8,7 @@ def kj_test( deps = [], tags = [], size = "medium", + local_defines = [], **kwargs): test_name = src.removesuffix(".c++") binary_name = test_name + "_binary" @@ -18,6 +20,11 @@ def kj_test( "@capnp-cpp//src/kj:kj-test", "//build/deps:linkopts_default", ] + deps, + # Under --//:io_backend=rust the native kj::setupAsyncIo()/UnixEventPort aren't linked, so + # tests that need an event loop swap to the tokio backend behind + # `#if WORKERD_RUST_IO_BACKEND_RUST`. Supplied to every kj_test TU here (cheap, harmless + # where unused) instead of repeating it per target. + local_defines = local_defines + rust_io_backend_local_defines(), linkstatic = select({ "@platforms//os:linux": 0, "//conditions:default": 1, diff --git a/build/rust_io_backend.bzl b/build/rust_io_backend.bzl new file mode 100644 index 00000000000..3aafc672f8c --- /dev/null +++ b/build/rust_io_backend.bzl @@ -0,0 +1,165 @@ +"""Build-graph hermeticity check + shared toggles for workerd's Rust I/O backend. + +An aspect walks the transitive dependency graph of a guarded target and, in the rust I/O +config, fails ANALYSIS if the graph reaches a forbidden concrete C++ I/O library (the ones +the rust layer replaces), printing one example dependency path to each offender. + +Why a graph assertion and not a link check: if kj-async-os sneaks back into the rust-config +graph, the failure mode is a DOUBLE-DEFINITION (kj::setupAsyncIo and kj::UnixEventPort's +members are defined by both the native TUs and the tokio shim), and with static archives the +winner is link-order-dependent -- possibly a loud duplicate-symbol error, possibly the wrong +event loop silently winning. Undefined-symbol errors only backstop the opposite (absent-lib) +direction. This aspect catches the double-definition direction deterministically, at analysis. + +IMPORTANT -- the gate is opt-in: //src/workerd/server:rust-io-hermeticity is tagged `manual`, +so `bazel build //...` never runs it. The rust-config CI lane must build it EXPLICITLY for +the guarantee to hold. + +The forbidden set grows in lockstep with what `=rust` means per migration stage (v1: the +tokio event loop + I/O; later stages append kj-http/kj-tls, then capnp-rpc). Deliberately +ALLOWED today: @capnp-cpp//src/kj:kj-async-core and :kj-async-io (the abstract Promise and +stream/Network layers the rust backend itself is built on), kj-http/kj-tls/capnp-rpc (still +the only implementation in both configs; they consume abstract streams, no OS I/O), and the +kj-gzip/kj-brotli codecs (not a transport). +""" + +visibility("public") + +def rust_io_backend_local_defines(): + """local_defines for TUs that `#if WORKERD_RUST_IO_BACKEND_RUST` (the declared seam points). + + Kept per-target rather than a repo-global define so the seam stays enumerable: the only + places allowed to diverge per backend are the targets that ask for this. + """ + return select({ + "//:io_backend_rust": ["WORKERD_RUST_IO_BACKEND_RUST=1"], + "//conditions:default": [], + }) + +# Forbidden concrete C++ I/O targets, as "//package:target" label suffixes (suffix-matched so +# bzlmod repo-name canonicalization doesn't have to be spelled out). The single source of truth. +_FORBIDDEN = [ + # The kj OS event loop / socket layer (setupAsyncIo, UnixEventPort/Win32IocpEventPort), + # replaced by kj-rs-tokio + kj-rs-io. + "//src/kj:kj-async-os", +] + +RustIoForbiddenInfo = provider( + doc = "One example dependency path to each forbidden C++ I/O target a subgraph reaches.", + fields = { + "paths": "dict of forbidden-label-suffix -> example path (list of label strings)", + }, +) + +def _forbidden_suffix(label_str): + for suffix in _FORBIDDEN: + if label_str.endswith(suffix): + return suffix + return None + +def _rust_io_forbidden_aspect_impl(target, ctx): + label_str = str(target.label) + paths = {} + + hit = _forbidden_suffix(label_str) + if hit != None: + paths[hit] = [label_str] + + # Scan ctx.rule.attr generically (attr_aspects = ["*"]) rather than enumerating + # deps/implementation_deps/etc. per rule kind. + for attr_name in dir(ctx.rule.attr): + value = getattr(ctx.rule.attr, attr_name) + dep_targets = [] + if type(value) == "list": + for item in value: + if type(item) == "Target": + dep_targets.append(item) + elif type(value) == "Target": + dep_targets.append(value) + + for dep in dep_targets: + if RustIoForbiddenInfo in dep: + for forbidden, subpath in dep[RustIoForbiddenInfo].paths.items(): + if forbidden not in paths: + paths[forbidden] = [label_str] + subpath + + return [RustIoForbiddenInfo(paths = paths)] + +rust_io_forbidden_aspect = aspect( + implementation = _rust_io_forbidden_aspect_impl, + attr_aspects = ["*"], + doc = "Propagates RustIoForbiddenInfo up the dependency graph.", +) + +def _rust_io_hermeticity_impl(ctx): + info = ctx.attr.target[RustIoForbiddenInfo] + paths = info.paths + + report = ctx.actions.declare_file(ctx.label.name + ".txt") + + if ctx.attr.enforce and len(paths) > 0: + lines = [ + "", + "Rust-I/O hermeticity FAILED for {} in the rust I/O config.".format( + str(ctx.attr.target.label), + ), + "", + "Its transitive dependency graph reaches {} forbidden concrete C++ ".format(len(paths)) + + "I/O target(s) that the rust I/O layer (the tokio event loop + tokio", + "sockets) is meant to keep off the build. Each is a real migration item; the", + "example dependency edge shows one path that pulls it in:", + "", + ] + for forbidden in sorted(paths.keys()): + path = paths[forbidden] + lines.append(" [FORBIDDEN] {}".format(forbidden)) + lines.append(" reached via:") + for i, hop in enumerate(path): + lines.append(" {}{}".format(" " * i, hop)) + lines.append("") + lines.append( + "Fix by removing the offending edge (migrate the code to the tokio-backed I/O layer,", + ) + lines.append( + "or drop the forbidden dep from that target's deps under select(io_backend_rust)),", + ) + lines.append( + "or -- if this is a deliberate remaining kj-mode/in-process site -- adjust the", + ) + lines.append( + "forbidden set in build/rust_io_backend.bzl (a reviewable change).", + ) + fail("\n".join(lines)) + + # Not enforcing (cxx config), or clean: emit a report artifact and a summary. + if len(paths) == 0: + summary = "rust-io-hermeticity: OK -- {} reaches 0 forbidden C++ I/O targets.".format( + str(ctx.attr.target.label), + ) + else: + summary = ("rust-io-hermeticity: {} reaches {} forbidden C++ I/O target(s) " + + "(NOT enforced in this config; --//:io_backend=rust would fail): {}").format( + str(ctx.attr.target.label), + len(paths), + ", ".join(sorted(paths.keys())), + ) + ctx.actions.write(report, summary + "\n") + return [DefaultInfo(files = depset([report]))] + +rust_io_hermeticity = rule( + implementation = _rust_io_hermeticity_impl, + doc = "Fails analysis (when enforced) if `target` transitively depends on a forbidden " + + "concrete C++ I/O library, naming the offending edge.", + attrs = { + "target": attr.label( + mandatory = True, + aspects = [rust_io_forbidden_aspect], + doc = "The target whose transitive deps are checked (e.g. the workerd binary).", + ), + "enforce": attr.bool( + default = False, + doc = "Whether reaching a forbidden target fails analysis. Set from a " + + "select() on //:io_backend_rust at the instantiation site.", + ), + }, +) diff --git a/deps/rust/Cargo.lock b/deps/rust/Cargo.lock index 26162790ba2..a37fe168c8d 100644 --- a/deps/rust/Cargo.lock +++ b/deps/rust/Cargo.lock @@ -444,6 +444,7 @@ dependencies = [ "ada-url", "anyhow", "async-trait", + "bytes", "capnp", "capnp-rpc", "capnpc", @@ -467,6 +468,7 @@ dependencies = [ "scratch", "serde", "serde_json", + "socket2", "static_assertions", "swc_common", "swc_ts_fast_strip", @@ -535,6 +537,16 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + [[package]] name = "fastrand" version = "2.5.0" @@ -1048,9 +1060,9 @@ dependencies = [ [[package]] name = "num-integer" -version = "0.1.46" +version = "0.1.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b" dependencies = [ "num-traits", ] @@ -1532,6 +1544,16 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + [[package]] name = "simd-adler32" version = "0.3.10" @@ -2025,13 +2047,27 @@ version = "1.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" dependencies = [ + "bytes", "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2", + "tokio-macros", "windows-sys", ] +[[package]] +name = "tokio-macros" +version = "2.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "tracing" version = "0.1.44" diff --git a/deps/rust/Cargo.toml b/deps/rust/Cargo.toml index 2c01d2b40c5..4249130a312 100644 --- a/deps/rust/Cargo.toml +++ b/deps/rust/Cargo.toml @@ -32,6 +32,7 @@ syn = { version = "2", features = ["full"] } ada-url = { version = "4", default-features = false, features = ["std"] } anyhow = "1" async-trait = { version = "0", default-features = false } +bytes = "1" capnp = "0" capnpc = "0" capnp-rpc = "0" @@ -48,9 +49,9 @@ ruff_python_parser = { git = "https://github.com/astral-sh/ruff", tag = "0.12.1" # param_extractor depends on unbounded_depth feature serde_json = { version = "1", features = ["unbounded_depth"] } serde = { version = "1", features = ["derive"] } +socket2 = "0.6" thiserror = "2" -# tokio is huge, let's enable only features when we actually need them. -tokio = { version = "1", default-features = false, features = ["net", "rt", "rt-multi-thread", "time"] } +tokio = { version = "1", default-features = false, features = ["io-util", "macros", "net", "rt", "rt-multi-thread", "signal", "sync", "time"] } tracing = { version = "0", default-features = false, features = ["std"] } swc_common = "25" swc_ts_fast_strip = "57" diff --git a/src/rust/cxx/AGENTS.md b/src/rust/cxx/AGENTS.md index bb8fb9d3ca2..49ab1fd4b2c 100644 --- a/src/rust/cxx/AGENTS.md +++ b/src/rust/cxx/AGENTS.md @@ -33,9 +33,29 @@ Bazel module, Cargo workspace, toolchain configuration, or external `workerd-cxx - `src/` and `include/` — cxx Rust and C++ runtimes - `syntax/`, `gen/`, and `macro/` — bridge parser and code generators - `kj-rs/` — KJ promises/futures, exceptions, ownership, refcounting, dates, and `Maybe` +- `kj-rs-tokio/` — `TokioEventPort`: a `kj::EventPort` backed by a per-thread tokio + `current_thread` runtime, plus `setupTokioAsyncIo()` (no I/O providers) and + `kj_rs_tokio::spawn()` +- `kj-rs-io/` — KJ async I/O interfaces over tokio (`kj::AsyncIoStream`, listeners, + `kj::Network`, providers, signals, file watching), `kj_rs_io::setupTokioAsyncIo()` as the + drop-in `kj::setupAsyncIo()` replacement, the stream unwrap fast path, and + `serve_kj_stream()` for Rust servers consuming KJ streams - `tests/` and `kj-rs/tests/` — Rust and C++ bridge integration tests - `tools/bazel/` — Bazel bridge-generation macro used by this component's tests +## Async bridge semantics + +- Marking a fn `async` in `extern "Rust"` yields a `kj::Promise` in C++; `async` in + `extern "C++"` yields an `impl Future` in Rust. +- Bridged `kj::Promise`s are **eager by default**: the Rust future is polled to its first + suspension point at the call (KJ code assumes hot promises), so callers never need + `.eagerlyEvaluate(nullptr)`. `RustFuture::lazily()` (kj-rs/future.h) is the C++-side + escape hatch for the rare cold case. +- The waker bridge is single-threaded: a Rust `.await` of a KJ promise links to the + `FuturePollEvent` via an intrusive weak link (`RustPromiseAwaiter::link` / + `FuturePollEvent::leaves`), and a cloned waker is a same-thread `FutureWakerCell` that arms + the `FuturePollEvent` directly (no atomics, no cross-thread fulfiller). + ## Conventions - Follow the parent `src/rust/AGENTS.md` and repository `AGENTS.md`. diff --git a/src/rust/cxx/kj-rs-io/BUILD.bazel b/src/rust/cxx/kj-rs-io/BUILD.bazel new file mode 100644 index 00000000000..cceafee0af8 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/BUILD.bazel @@ -0,0 +1,79 @@ +load("@rules_rust//rust:defs.bzl", "rust_library", "rust_test") +load("//:build/wd_cc_library.bzl", "wd_cc_library") +load("//src/rust/cxx/tools/bazel:rust_cxx_bridge.bzl", "rust_cxx_bridge") + +wd_cc_library( + name = "kj-rs-io-lib", + srcs = glob(["*.c++"]), + hdrs = glob(["*.h"]), + include_prefix = "kj-rs-io", + linkstatic = select({ + "@platforms//os:windows": True, + "//conditions:default": False, + }), + strip_include_prefix = "/src/rust/cxx/kj-rs-io", + visibility = ["//visibility:public"], + deps = [ + ":bridge", + "//src/rust/cxx/kj-rs-tokio:kj-rs-tokio-lib", + ], +) + +rust_library( + name = "kj-rs-io", + srcs = glob(["*.rs"]), + compile_data = glob(["*.h"]), + edition = "2024", + link_deps = [ + ":bridge", + ":kj-rs-io-lib", + ], + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + visibility = ["//visibility:public"], + deps = [ + "//src/rust/cxx", + "//src/rust/cxx/kj-rs", + "//src/rust/cxx/kj-rs-tokio", + "@crates_vendor//:bytes", + "@crates_vendor//:libc", + "@crates_vendor//:socket2", + "@crates_vendor//:tokio", + ], +) + +rust_test( + name = "kj-rs-io_test", + crate = "kj-rs-io", + edition = "2024", + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), +) + +rust_cxx_bridge( + name = "bridge", + src = "ffi.rs", + hdrs = ["unwrap.h"], + include_prefix = "kj-rs-io", + visibility = ["//visibility:public"], + deps = [ + "//src/rust/cxx/kj-rs", + "@capnp-cpp//src/kj:kj", + # kj-rs-io IS the kj-side implementation of the rust I/O backend: it derives from the + # abstract kj::AsyncIoStream / kj::Network / kj::LowLevelAsyncIoProvider interfaces and + # reuses portable helpers (kj::newOneWayPipe, kj::CidrRange, the AsyncInputStream / + # AsyncOutputStream default-method vtable slots). Those live in the ABSTRACT layers + # :kj-async-core (Promise machinery) + :kj-async-io (abstract streams / Network / CIDR; + # no event loop). It does NOT call kj::setupAsyncIo / kj::UnixEventPort, so it never + # needs :kj-async-os (the concrete OS event loop) -- the tokio-backed EventPort + # (kj-rs-tokio) replaces it. A stray setupAsyncIo would be an undefined-symbol link + # error, the intended backstop. + "@capnp-cpp//src/kj:kj-async-core", + "@capnp-cpp//src/kj:kj-async-io", + "//src/rust/cxx:core", + ], +) diff --git a/src/rust/cxx/kj-rs-io/async-io.c++ b/src/rust/cxx/kj-rs-io/async-io.c++ new file mode 100644 index 00000000000..71638c10269 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/async-io.c++ @@ -0,0 +1,531 @@ +#include "kj-rs-io/async-io.h" + +#include + +#include + +#if _WIN32 +#include +#else +#include +#include +#include +#include +#if __APPLE__ || __FreeBSD__ || __OpenBSD__ || __NetBSD__ || __DragonFly__ +#include +#endif +#endif + +namespace kj_rs_io { + +// ======================================================================================= +// TokioAsyncIoStream + +kj::Promise TokioAsyncIoStream::tryRead(void *buffer, size_t minBytes, size_t maxBytes) { + return stream_try_read( + *inner, ::rust::Slice(reinterpret_cast(buffer), maxBytes), minBytes); +} + +kj::Promise TokioAsyncIoStream::write(kj::ArrayPtr buffer) { + return stream_write(*inner, ::rust::Slice(buffer.begin(), buffer.size())); +} + +kj::Promise TokioAsyncIoStream::write( + kj::ArrayPtr> pieces) { + // Sequential write-all of each piece (each single write is eager-by-default, preserving + // hot-write semantics for the whole sequence). + // TODO(perf): vectored writes via try_write_vectored. + return writePieces(pieces); +} + +kj::Promise TokioAsyncIoStream::writePieces( + kj::ArrayPtr> pieces) { + for (auto piece: pieces) { + co_await write(piece); + } +} + +kj::Promise TokioAsyncIoStream::whenWriteDisconnected() { + return stream_when_write_disconnected(*inner); +} + +void TokioAsyncIoStream::shutdownWrite() { + stream_shutdown_write(*inner); +} + +void TokioAsyncIoStream::getsockopt(int level, int option, void *value, kj::uint *length) { + // The platform seam lives on the Rust side (stream_getsockopt); errors surface as kj + // exceptions, like KJ_SYSCALL. Raw socklen in/out semantics: the syscall's reported length + // is mirrored back verbatim. + *length = stream_getsockopt( + *inner, level, option, ::rust::Slice(reinterpret_cast(value), *length)); +} + +void TokioAsyncIoStream::setsockopt(int level, int option, const void *value, kj::uint length) { + stream_setsockopt(*inner, level, option, + ::rust::Slice(reinterpret_cast(value), length)); +} + +void TokioAsyncIoStream::getsockname(struct sockaddr *addr, kj::uint *length) { + auto bytes = stream_local_addr(*inner); + // Mirror the raw syscall's truncation semantics: copy what fits into the caller's buffer, + // report the address's full length. + memcpy(addr, bytes.data(), kj::min(bytes.size(), *length)); + *length = bytes.size(); +} + +void TokioAsyncIoStream::getpeername(struct sockaddr *addr, kj::uint *length) { + auto bytes = stream_peer_addr(*inner); + memcpy(addr, bytes.data(), kj::min(bytes.size(), *length)); + *length = bytes.size(); +} + +kj::Maybe TokioAsyncIoStream::getFd() const { +#if _WIN32 + // On Windows the underlying handle is a winsock SOCKET, not a Unix fd; it is exposed via + // getWin32Handle() below instead. (Validated by Windows CI.) + return kj::none; +#else + int64_t handle = -1; + if (kj::runCatchingExceptions([&]() { handle = stream_raw_handle(*inner); }) == kj::none) { + // On unix the raw socket handle is the fd, widened losslessly to int64 by the bridge. + return static_cast(handle); + } + return kj::none; +#endif +} + +#if _WIN32 +// Validated by Windows CI; mirrors the unix getFd() arm (and kj's own win32 AsyncStreamFd, +// which returns its SOCKET cast to void* -- capnproto async-io-win32.c++). +kj::Maybe TokioAsyncIoStream::getWin32Handle() const { + int64_t handle = -1; + if (kj::runCatchingExceptions([&]() { handle = stream_raw_handle(*inner); }) == kj::none) { + return reinterpret_cast(static_cast(handle)); + } + return kj::none; +} +#endif + +::rust::Box unwrapTokioStream(kj::AsyncIoStream &stream) { + KJ_IF_SOME(tokioStream, kj::dynamicDowncastIfAvailable(stream)) { + return tokioStream.unwrap(); + } + KJ_FAIL_REQUIRE("stream is not a kj-rs-io tokio-backed stream; cannot unwrap"); +} + +// ======================================================================================= +// TokioConnectionReceiver + +namespace { + +// Builds the accepted connection's kj::PeerIdentity, mirroring KJ's SocketAddress::getIdentity() +// (kj/async-io-unix.c++): NetworkPeerIdentity wrapping the peer's address for TCP peers (its +// toString() is "ip:port" / "[v6]:port", byte-identical to KJ's format -- workerd's HTTP +// listener puts this string in the cf blob's clientIp), LocalPeerIdentity with the peer's +// process credentials for unix sockets, UnknownPeerIdentity otherwise. +kj::Own peerIdentityFromSockaddr( + struct sockaddr *sa, kj::uint addrlen, [[maybe_unused]] kj::AsyncIoStream &stream) { + switch (sa->sa_family) { + case AF_INET: + case AF_INET6: { + // The identity's NetworkAddress uses an allow-all filter (not the listener's): it exists + // for toString()/getAddress(); restrictPeers enforcement on this listener already happened + // in the accept loop. (KJ instead threads the listener's filter through, which only + // matters if a caller connect()s back through the identity address -- see PeerFilter's + // immobility note for why we don't hold a reference to a possibly-narrower filter here.) + static PeerFilter allowAll; + auto address = network_get_sockaddr( + ::rust::Slice(reinterpret_cast(sa), addrlen)); + return kj::NetworkPeerIdentity::newInstance( + kj::heap(kj::mv(address), allowAll)); + } +#if !_WIN32 + case AF_UNIX: { + // Same credential sources and invalid-value handling as KJ (SO_PEERCRED on Linux, + // LOCAL_PEERCRED/LOCAL_PEERPID on BSDs/macOS; OpenBSD defines SO_PEERCRED but with a + // different interface, so it uses the LOCAL_PEERCRED arm). + kj::LocalPeerIdentity::Credentials result; +#if defined(SO_PEERCRED) && !__OpenBSD__ + struct ucred creds; + kj::uint length = sizeof(creds); + stream.getsockopt(SOL_SOCKET, SO_PEERCRED, &creds, &length); + if (creds.pid > 0) { + result.pid = creds.pid; + } + if (creds.uid != static_cast(-1)) { + result.uid = creds.uid; + } +#elifdef LOCAL_PEERCRED + struct xucred creds; + kj::uint length = sizeof(creds); + stream.getsockopt(SOL_LOCAL, LOCAL_PEERCRED, &creds, &length); + KJ_ASSERT(length == sizeof(creds)); + if (creds.cr_uid != static_cast(-1)) { + result.uid = creds.cr_uid; + } +#ifdef LOCAL_PEERPID + pid_t pid; + length = sizeof(pid); + stream.getsockopt(SOL_LOCAL, LOCAL_PEERPID, &pid, &length); + KJ_ASSERT(length == sizeof(pid)); + if (pid > 0) { + result.pid = pid; + } +#endif +#endif + return kj::LocalPeerIdentity::newInstance(result); + } +#endif // !_WIN32 + default: + return kj::UnknownPeerIdentity::newInstance(); + } +} + +} // namespace + +kj::Promise> TokioConnectionReceiver::accept() { + return acceptImpl(false).then( + [](kj::AuthenticatedStream authenticated) { return kj::mv(authenticated.stream); }); +} + +kj::Promise TokioConnectionReceiver::acceptAuthenticated() { + return acceptImpl(true); +} + +kj::Promise TokioConnectionReceiver::acceptImpl(bool authenticated) { + for (;;) { + auto stream = co_await listener_accept(*inner); + // restrictPeers / NetworkFilter enforcement, mirroring KJ: a connection from a disallowed + // peer is silently dropped and we keep accepting. + struct sockaddr_storage addr; + memset(&addr, 0, sizeof(addr)); + kj::uint addrlen = 0; + KJ_IF_SOME(exception, kj::runCatchingExceptions([&]() { + auto bytes = stream_peer_addr(*stream); + KJ_ASSERT(bytes.size() <= sizeof(addr), "sockaddr too large"); + memcpy(&addr, bytes.data(), bytes.size()); + addrlen = bytes.size(); + })) { + // The peer can reset the connection between tokio's accept() and this call, in which + // case getpeername fails (EINVAL on macOS, ENOTCONN elsewhere). The connection is dead; + // drop it and keep accepting. This must NOT throw: an exception here propagates out of + // the server's accept loop and takes down the whole process (observed as a fatal + // uncaught kj::Exception under client abort storms). Unlike KJ's native accept path, + // which gets the peer address atomically from accept4(), we re-derive it and so must + // tolerate the race. Log at INFO (off by default) for observability under abort storms. + KJ_LOG(INFO, "dropping accepted connection; could not read peer address", exception); + continue; + } + if (!filter.shouldAllow(reinterpret_cast(&addr), addrlen)) { + // Drop the disallowed connection and wait for the next one. + continue; + } + kj::AuthenticatedStream result; + result.stream = kj::heap(kj::mv(stream)); + if (authenticated) { + result.peerIdentity = peerIdentityFromSockaddr( + reinterpret_cast(&addr), addrlen, *result.stream); + } else { + result.peerIdentity = kj::UnknownPeerIdentity::newInstance(); + } + co_return kj::mv(result); + } +} + +kj::uint TokioConnectionReceiver::getPort() { + return listener_port(*inner); +} + +void TokioConnectionReceiver::getsockopt(int level, int option, void *value, kj::uint *length) { + *length = listener_getsockopt( + *inner, level, option, ::rust::Slice(reinterpret_cast(value), *length)); +} + +void TokioConnectionReceiver::setsockopt( + int level, int option, const void *value, kj::uint length) { + listener_setsockopt(*inner, level, option, + ::rust::Slice(reinterpret_cast(value), length)); +} + +void TokioConnectionReceiver::getsockname(struct sockaddr *addr, kj::uint *length) { + auto bytes = listener_local_addr(*inner); + // Mirror the raw syscall's truncation semantics (see TokioAsyncIoStream::getsockname). + memcpy(addr, bytes.data(), kj::min(bytes.size(), *length)); + *length = bytes.size(); +} + +// ======================================================================================= +// TokioNetworkAddress / TokioNetwork + +kj::Promise> TokioNetworkAddress::connect() { + // KJ contract (NetworkAddressImpl::connect() in kj/async-io-unix.c++): callers may drop the + // NetworkAddress while the returned promise is still pending. We honor this by cloning the + // resolved address list into a coroutine frame local: the frame owns the copy, so it + // survives every co_await and is dropped on completion/cancellation. connect() may be called + // repeatedly on the same address, so we clone rather than move `*inner` out of `this`. + // `filter` is the network's filter (provider-owned, shared) and must outlive the promise, + // per KJ (where it lives in the long-lived provider). + auto addr = address_clone(*inner); + size_t count = address_count(*addr); + KJ_REQUIRE(count > 0, "no addresses to connect to"); + + // Try each resolved address in order; a filter block or connect error falls through to the + // next one, and the last address's exception propagates (KJ parity). + kj::Maybe lastException; + for (size_t i = 0; i < count; i++) { + kj::Maybe> stream; + try { + auto raw = address_raw_sockaddr(*addr, i); + // Copy into sockaddr_storage for alignment (rust::Vec data is 1-aligned). + struct sockaddr_storage storage; + memset(&storage, 0, sizeof(storage)); + KJ_REQUIRE(raw.size() <= sizeof(storage), "sockaddr too large"); + memcpy(&storage, raw.data(), raw.size()); + if (!filter.shouldAllow(reinterpret_cast(&storage), raw.size())) { + // Exact KJ error text; error-string parity matters to callers. + lastException = KJ_EXCEPTION(FAILED, "connect() blocked by restrictPeers()"); + } else { + stream = kj::heap(co_await address_connect_index(*addr, i)); + } + } catch (...) { + // Note: getCaughtExceptionAsKj() rethrows kj::CanceledException, so cancellation still + // propagates out of this coroutine instead of being folded into lastException. + lastException = kj::getCaughtExceptionAsKj(); + } + KJ_IF_SOME(s, stream) { + co_return kj::mv(s); + } + } + + kj::throwFatalException(kj::mv(KJ_ASSERT_NONNULL(lastException))); +} + +kj::Own TokioNetworkAddress::listen() { + return kj::heap(address_listen(*inner), filter); +} + +kj::Own TokioNetworkAddress::clone() { + return kj::heap(address_clone(*inner), filter); +} + +kj::String TokioNetworkAddress::toString() { + auto text = address_to_string(*inner); + return kj::heapString(text.data(), text.size()); +} + +kj::Promise> TokioNetwork::parseAddress( + kj::StringPtr addr, kj::uint portHint) { + KJ_REQUIRE(portHint < 65536, "port hint too large", portHint); + // The Rust side takes an owned copy: the caller's buffer need not outlive this call. + // + // Note: unlike KJ, disallowed (restrictPeers) DNS results are not dropped here; they are + // rejected at connect()/accept() time instead. See TokioNetworkAddress. + return network_parse_address( + ::rust::String(addr.begin(), addr.size()), static_cast(portHint)) + .then([this](::rust::Box address) -> kj::Own { + return kj::heap(kj::mv(address), filter); + }); +} + +kj::Own TokioNetwork::getSockaddr(const void *sockaddr, kj::uint len) { + // KJ parity: getSockaddr() rejects filtered addresses eagerly (same error text as KJ). + KJ_REQUIRE(filter.shouldAllow(reinterpret_cast(sockaddr), len), + "address blocked by restrictPeers()"); + return kj::heap(network_get_sockaddr(::rust::Slice( + reinterpret_cast(sockaddr), len)), + filter); +} + +kj::Own TokioNetwork::restrictPeers( + kj::ArrayPtr allow, kj::ArrayPtr deny) { + // The child references this network's filter chain: this network must outlive the returned + // one (same constraint as KJ's networks). + return kj::heap(*this, allow, deny); +} + +// ======================================================================================= +// TokioLowLevelAsyncIoProvider + +namespace { + +#if _WIN32 +// Normalizes KJ's fd-wrapping flags so Rust always receives a SOCKET it owns, in non-blocking +// mode: the windows arm of prepareFd, mirroring the unix arm below in *effect*, not mechanism. +// kj's own win32 provider (capnproto async-io-win32.c++: OwnedFd, NEW_FD_FLAGS) never dups a +// borrowed socket -- it merely skips closesocket() on destruction when TAKE_OWNERSHIP is +// absent -- ignores ALREADY_CLOEXEC entirely (there is no CLOEXEC on Windows; handle +// inheritance is the analogue), and never toggles non-blocking mode (it uses overlapped I/O, +// not readiness). Rust's OwnedSocket has no "don't close" mode, so a borrowed socket is +// duplicated (WSADuplicateSocketW + WSASocketW, non-inheritable) into a handle Rust can own; +// and tokio/mio's readiness model requires non-blocking sockets, so FIONBIO is set unless the +// caller declared ALREADY_NONBLOCK (the duplicate shares the underlying socket state, so this +// is observed through the caller's handle too, matching the unix dup()+O_NONBLOCK behavior). +// Validated by Windows CI. +uintptr_t prepareFd(uintptr_t fd, kj::uint flags) { + SOCKET sock = static_cast(fd); + if ((flags & kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP) == 0) { + WSAPROTOCOL_INFOW info; + KJ_WINSOCK(WSADuplicateSocketW(sock, GetCurrentProcessId(), &info)); + SOCKET duped = WSASocketW(FROM_PROTOCOL_INFO, FROM_PROTOCOL_INFO, FROM_PROTOCOL_INFO, &info, 0, + WSA_FLAG_OVERLAPPED | WSA_FLAG_NO_HANDLE_INHERIT); + if (duped == INVALID_SOCKET) { + KJ_FAIL_WIN32("WSASocketW()", WSAGetLastError()); + } + sock = duped; + } + // ALREADY_NONBLOCK does not exist on Windows (kj declares it under #if !_WIN32), so callers + // cannot assert pre-set non-blocking mode; always enable it (idempotent). + u_long mode = 1; + KJ_WINSOCK(ioctlsocket(sock, FIONBIO, &mode)); + return static_cast(sock); +} +#else +// Normalizes KJ's fd-wrapping flags so Rust always receives an fd it owns, with CLOEXEC set and +// in non-blocking mode. +int prepareFd(int fd, kj::uint flags) { + if ((flags & kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP) == 0) { + // dup() shares the open file description — the O_NONBLOCK set below is observed through the + // caller's fd too, matching KJ (which sets O_NONBLOCK on the caller's fd directly) — while + // giving Rust a descriptor it can own and close. + int duped; + KJ_SYSCALL(duped = ::dup(fd)); + fd = duped; + KJ_SYSCALL(fcntl(fd, F_SETFD, FD_CLOEXEC)); + } else if ((flags & kj::LowLevelAsyncIoProvider::ALREADY_CLOEXEC) == 0) { + KJ_SYSCALL(fcntl(fd, F_SETFD, FD_CLOEXEC)); + } + if ((flags & kj::LowLevelAsyncIoProvider::ALREADY_NONBLOCK) == 0) { + int fl; + KJ_SYSCALL(fl = fcntl(fd, F_GETFL)); + if ((fl & O_NONBLOCK) == 0) { + KJ_SYSCALL(fcntl(fd, F_SETFL, fl | O_NONBLOCK)); + } + } + return fd; +} + +// kj::AsyncInputStream over an arbitrary readable fd (pipe, socket, character device). +class TokioInputStreamFd final: public kj::AsyncInputStream { + public: + explicit TokioInputStreamFd(::rust::Box inner): inner(kj::mv(inner)) {} + + kj::Promise tryRead(void *buffer, size_t minBytes, size_t maxBytes) override { + return input_fd_try_read( + *inner, ::rust::Slice(reinterpret_cast(buffer), maxBytes), minBytes); + } + + private: + ::rust::Box inner; +}; + +// kj::AsyncOutputStream over an arbitrary writable fd. +class TokioOutputStreamFd final: public kj::AsyncOutputStream { + public: + explicit TokioOutputStreamFd(::rust::Box inner): inner(kj::mv(inner)) {} + + kj::Promise write(kj::ArrayPtr buffer) override { + return output_fd_write(*inner, ::rust::Slice(buffer.begin(), buffer.size())); + } + + kj::Promise write(kj::ArrayPtr> pieces) override { + for (auto piece: pieces) { + co_await write(piece); + } + } + + // Pipes/arbitrary fds have no portable disconnect detection here; KJ allows a never-resolving + // promise for such streams. + kj::Promise whenWriteDisconnected() override { + return kj::NEVER_DONE; + } + + private: + ::rust::Box inner; +}; +#endif // _WIN32 (prepareFd platform arms; the pipe-fd stream classes above are unix-only) + +} // namespace + +kj::Own TokioLowLevelAsyncIoProvider::wrapInputFd(Fd fd, kj::uint flags) { +#if _WIN32 + // KJ parity: on win32, LowLevelAsyncIoProvider::Fd is documented as a SOCKET (async-io.h: + // "On Windows, the `fd` parameter to each of these methods must be a SOCKET"), and kj's own + // win32 provider implements wrapInputFd/wrapOutputFd *identically* to wrapSocketFd + // (async-io-win32.c++: all three wrap the SOCKET in AsyncStreamFd) -- even kj's "pipes" on + // Windows are loopback-TCP socketpairs (newOsSocketpair). So there is no pipe-HANDLE tier to + // implement; delegate to the tested socket path. Validated by Windows CI. + return kj::heap(wrap_socket_fd(static_cast(prepareFd(fd, flags)))); +#else + return kj::heap(wrap_input_fd(prepareFd(fd, flags))); +#endif +} + +kj::Own TokioLowLevelAsyncIoProvider::wrapOutputFd(Fd fd, kj::uint flags) { +#if _WIN32 + // See wrapInputFd above: win32 Fd is a SOCKET and kj's win32 wrapOutputFd == wrapSocketFd. + // Validated by Windows CI. + return kj::heap(wrap_socket_fd(static_cast(prepareFd(fd, flags)))); +#else + return kj::heap(wrap_output_fd(prepareFd(fd, flags))); +#endif +} + +kj::Own TokioLowLevelAsyncIoProvider::wrapSocketFd(Fd fd, kj::uint flags) { + // `Fd` is int on unix and uintptr_t (SOCKET) on win32; the bridge carries it widened to + // int64 (a "raw socket handle") either way, with -1 == INVALID_SOCKET as the one sentinel. + return kj::heap(wrap_socket_fd(static_cast(prepareFd(fd, flags)))); +} + +kj::Promise> TokioLowLevelAsyncIoProvider::wrapConnectingSocketFd( + Fd fd, const struct sockaddr *addr, kj::uint addrlen, kj::uint flags) { + int64_t prepared = static_cast(prepareFd(fd, flags)); + // The Rust side takes an owned copy of the sockaddr: the caller's pointer need not outlive + // this call (KJ's own implementation copies too). + ::rust::Vec addrCopy; + addrCopy.reserve(addrlen); + const uint8_t *addrBytes = reinterpret_cast(addr); + for (kj::uint i = 0; i < addrlen; i++) { + addrCopy.push_back(addrBytes[i]); + } + return wrap_connecting_socket_fd(prepared, kj::mv(addrCopy)) + .then([](::rust::Box stream) -> kj::Own { + return kj::heap(kj::mv(stream)); + }); +} + +kj::Own TokioLowLevelAsyncIoProvider::wrapListenSocketFd( + Fd fd, NetworkFilter &filter, kj::uint flags) { + // `filter` applies to accepted connections (KJ parity); it must outlive the receiver. + return kj::heap( + wrap_listen_fd(static_cast(prepareFd(fd, flags))), filter); +} + +// ======================================================================================= +// TokioAsyncIoProvider / setup + +kj::AsyncIoProvider::PipeThread TokioAsyncIoProvider::newPipeThread( + kj::Function startFunc) { + KJ_UNIMPLEMENTED("kj-rs-io does not implement newPipeThread() (workerd does not use it)"); +} + +TokioAsyncIoContext setupTokioAsyncIo() { + auto port = kj::heap(); + auto loop = kj::heap(*port); + auto waitScope = kj::heap(*loop); + auto lowLevelProvider = kj::heap(port->getTimer()); + auto provider = kj::heap(port->getTimer()); + return TokioAsyncIoContext{ + kj::mv(port), kj::mv(loop), kj::mv(waitScope), kj::mv(lowLevelProvider), kj::mv(provider)}; +} + +// ======================================================================================= +// Signals + +kj::Promise onSignal(int signum) { + // The bridged future is eager-by-default, so the tokio signal handler is registered as soon + // as the event loop runs, even if the caller parks the promise without awaiting it immediately. + return wait_for_signal(signum); +} + +} // namespace kj_rs_io diff --git a/src/rust/cxx/kj-rs-io/async-io.h b/src/rust/cxx/kj-rs-io/async-io.h new file mode 100644 index 00000000000..0418abd9b48 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/async-io.h @@ -0,0 +1,262 @@ +#pragma once +// kj-rs-io: tokio-backed implementations of KJ's async I/O interfaces. +// +// Everything here wraps an opaque Rust object (a native tokio TcpStream/UnixStream/TcpListener/ +// address list) and implements the corresponding KJ interface by calling `async fn`s across the +// cxx bridge, which return kj::Promises. Design points: +// +// - All promises must be awaited on the thread owning the kj_rs_tokio::TokioEventPort: the +// tokio I/O driver that delivers readiness for these sockets only runs while that KJ loop +// sleeps in the port's wait()/poll(). +// - Cancellation: dropping a returned kj::Promise drops the underlying Rust future, which +// releases the socket's readiness interest. A stream with a canceled read remains usable. +// - Unwrap fast path: every Rust-originated stream can be recovered as its native tokio object +// (see unwrapTokioStream), so Rust servers can serve a connection natively +// instead of crossing the FFI per read. Foreign kj streams are not unwrappable. +// +// Known stubs (all throw UNIMPLEMENTED, documented per method): newPipeThread(), capability +// streams (SCM_RIGHTS fd passing), datagram sockets, and named-service / abstract-unix-socket +// address forms. restrictPeers() IS implemented (via PeerFilter, a port of KJ's NetworkFilter); see +// TokioNetwork for enforcement points and the parse-time-filtering deviation. + +#include "kj-rs-io/ffi.rs.h" +#include "kj-rs-io/peer-filter.h" +#include "kj-rs-tokio/tokio-event-port.h" + +#include +#include + +namespace kj_rs_io { + +// A kj::AsyncIoStream backed by a native tokio TcpStream or UnixStream. +class TokioAsyncIoStream final: public kj::AsyncIoStream { + public: + explicit TokioAsyncIoStream(::rust::Box inner): inner(kj::mv(inner)) {} + + // AsyncInputStream. tryRead honors KJ's min-bytes contract: resolves with >= minBytes unless + // EOF is reached first (in which case the short count signals EOF). + kj::Promise tryRead(void *buffer, size_t minBytes, size_t maxBytes) override; + + // AsyncOutputStream. write() has write-all semantics; the multi-piece overload writes the + // pieces sequentially (no vectored-write optimization yet). + kj::Promise write(kj::ArrayPtr buffer) override; + kj::Promise write(kj::ArrayPtr> pieces) override; + + // Resolves when new writes are doomed (peer reset/hangup observed). Does not fire on a mere + // half-close (peer FIN), mirroring KJ. On non-Unix platforms the promise never resolves + // (KJ-on-Windows behavior). Safe to call multiple times concurrently. + kj::Promise whenWriteDisconnected() override; + + // AsyncIoStream. + void shutdownWrite() override; + void getsockopt(int level, int option, void *value, kj::uint *length) override; + void setsockopt(int level, int option, const void *value, kj::uint length) override; + void getsockname(struct sockaddr *addr, kj::uint *length) override; + void getpeername(struct sockaddr *addr, kj::uint *length) override; + kj::Maybe getFd() const override; +#if _WIN32 + // Validated by Windows CI; mirrors getFd(): on Windows the underlying socket is a winsock + // SOCKET, exposed as a void* handle (kj convention; getFd() returns none there). + kj::Maybe getWin32Handle() const override; +#endif + + // Unwrap fast path: moves the native tokio stream out, leaving this wrapper hollow (all + // further operations throw). No I/O promises may be in flight. Prefer the free function + // unwrapTokioStream() when holding only a kj::AsyncIoStream&. + ::rust::Box unwrap() { + return stream_take(*inner); + } + + private: + kj::Promise writePieces(kj::ArrayPtr> pieces); + + ::rust::Box inner; +}; + +// A kj::ConnectionReceiver backed by a native tokio TcpListener or UnixListener. Incoming +// connections from peers disallowed by `filter` (restrictPeers) are silently dropped and the +// accept loop continues, mirroring KJ. `filter` must outlive this receiver. +class TokioConnectionReceiver final: public kj::ConnectionReceiver { + public: + TokioConnectionReceiver( + ::rust::Box inner, kj::LowLevelAsyncIoProvider::NetworkFilter &filter) + : inner(kj::mv(inner)), + filter(filter) {} + + kj::Promise> accept() override; + kj::Promise acceptAuthenticated() override; + kj::uint getPort() override; + void getsockopt(int level, int option, void *value, kj::uint *length) override; + void setsockopt(int level, int option, const void *value, kj::uint length) override; + void getsockname(struct sockaddr *addr, kj::uint *length) override; + + private: + kj::Promise acceptImpl(bool authenticated); + + ::rust::Box inner; + kj::LowLevelAsyncIoProvider::NetworkFilter &filter; +}; + +// A kj::NetworkAddress holding pre-resolved socket addresses (DNS happens at parseAddress time, +// like KJ). connect() tries each address in order; listen() binds the first. +// +// connect() honors KJ's lifetime contract (NetworkAddressImpl::connect() in +// kj/async-io-unix.c++): the returned promise is a coroutine whose frame owns a copy of the +// resolved address list, so the caller may drop this NetworkAddress while the connect is still +// pending. +// +// `filter` is the restrictPeers filter of the kj::Network this address came from (allow-all +// for an unrestricted network) and must outlive this address and any promises it returns +// (in KJ the filter equally lives in the long-lived provider, so this matches upstream). +// Filtering is enforced at connect() time per address ("connect() blocked by +// restrictPeers()", KJ parity) and at accept() time on listeners; unlike KJ, disallowed DNS +// results are not already dropped at parse time (they fail at connect instead). +class TokioNetworkAddress final: public kj::NetworkAddress { + public: + TokioNetworkAddress(::rust::Box inner, PeerFilter &filter) + : inner(kj::mv(inner)), + filter(filter) {} + + kj::Promise> connect() override; + kj::Own listen() override; + kj::Own clone() override; + kj::String toString() override; + + private: + ::rust::Box inner; + PeerFilter &filter; +}; + +// The tokio-backed kj::Network. Supports the KJ address grammar subset workerd uses; see +// net.rs for the exact forms and documented deviations (no named services, no unix-abstract, +// no IPv6 scope IDs). +// +// restrictPeers() uses PeerFilter, a faithful port of KJ's NetworkFilter (semantics identical to +// kj::setupAsyncIo()'s networks). The returned network references this one's filter chain, so +// a network must outlive any networks derived from it via restrictPeers() (same constraint as +// KJ). Enforcement points: per-address connect()-time checks and accept()-time peer checks; +// parse-time DNS-result dropping is NOT implemented (blocked addresses fail at connect +// instead) -- see TokioNetworkAddress. +class TokioNetwork final: public kj::Network { + public: + TokioNetwork() = default; + TokioNetwork(TokioNetwork &parent, + kj::ArrayPtr allow, + kj::ArrayPtr deny) + : filter(allow, deny, parent.filter) {} + + kj::Promise> parseAddress( + kj::StringPtr addr, kj::uint portHint) override; + kj::Own getSockaddr(const void *sockaddr, kj::uint len) override; + kj::Own restrictPeers( + kj::ArrayPtr allow, kj::ArrayPtr deny) override; + + private: + // Default-constructed = allow everything (matches KJ's root networks). + PeerFilter filter; +}; + +// The tokio-backed kj::LowLevelAsyncIoProvider. Implements the socket-wrapping entry points on +// Unix and Windows (each wrap*Fd normalizes KJ's TAKE_OWNERSHIP/ALREADY_CLOEXEC/ALREADY_NONBLOCK +// flags, then hands an owned, non-blocking raw socket handle -- a Unix fd or a win32 SOCKET, +// widened to int64 -- to Rust). The pipe tier (wrapInputFd/wrapOutputFd) is Unix-only for now. +// wrapUnixSocketFd (capability streams) and wrapDatagramSocketFd keep their default-throwing +// implementations. +class TokioLowLevelAsyncIoProvider final: public kj::LowLevelAsyncIoProvider { + public: + explicit TokioLowLevelAsyncIoProvider(kj::Timer &timer): timer(timer) {} + + kj::Own wrapInputFd(Fd fd, kj::uint flags) override; + kj::Own wrapOutputFd(Fd fd, kj::uint flags) override; + kj::Own wrapSocketFd(Fd fd, kj::uint flags) override; + kj::Promise> wrapConnectingSocketFd( + Fd fd, const struct sockaddr *addr, kj::uint addrlen, kj::uint flags) override; + // `filter` applies to accepted connections (disallowed peers are dropped and the accept + // loop continues, like KJ); it must outlive the returned receiver. + kj::Own wrapListenSocketFd( + Fd fd, NetworkFilter &filter, kj::uint flags) override; + kj::Timer &getTimer() override { + return timer; + } + + private: + kj::Timer &timer; +}; + +// The tokio-backed kj::AsyncIoProvider. Pipes are KJ's in-memory pipes (port-agnostic, like +// kj::newOneWayPipe/newTwoWayPipe themselves); newPipeThread throws UNIMPLEMENTED (workerd does +// not use it); newCapabilityPipe keeps its default-throwing implementation. +class TokioAsyncIoProvider final: public kj::AsyncIoProvider { + public: + explicit TokioAsyncIoProvider(kj::Timer &timer): timer(timer) {} + + kj::OneWayPipe newOneWayPipe() override { + return kj::newOneWayPipe(); + } + kj::TwoWayPipe newTwoWayPipe() override { + return kj::newTwoWayPipe(); + } + kj::Network &getNetwork() override { + return network; + } + PipeThread newPipeThread( + kj::Function startFunc) + override; + kj::Timer &getTimer() override { + return timer; + } + + private: + TokioNetwork network; + kj::Timer &timer; +}; + +// Mirror of kj::AsyncIoContext (kj/async-io.h) for the tokio-backed loop: a drop-in replacement +// for kj::setupAsyncIo() at workerd.c++:1570. Additionally owns the event port / loop / +// WaitScope (which kj::setupAsyncIo keeps in thread-locals). +struct TokioAsyncIoContext { + // Destroyed in reverse declaration order: providers first (their rust::Boxes drop while the + // runtime still exists), then waitScope, then loop (asserts its queue is empty), then port + // (dropping the tokio runtime, canceling still-pending spawned tasks). I/O objects created + // *through* the providers (streams, listeners, addresses) must be destroyed before the + // context, as with kj::setupAsyncIo(). + kj::Own port; + kj::Own loop; + kj::Own waitScope; + kj::Own lowLevelProvider; + kj::Own provider; + + kj_rs_tokio::TokioEventPort &getPort() { + return *port; + } + kj::WaitScope &getWaitScope() { + return *waitScope; + } + kj::Timer &getTimer() { + return port->getTimer(); + } + kj::Network &getNetwork() { + return provider->getNetwork(); + } + kj::AsyncIoProvider &getProvider() { + return *provider; + } + kj::LowLevelAsyncIoProvider &getLowLevelProvider() { + return *lowLevelProvider; + } +}; + +// Sets up the current thread with a tokio-driven KJ event loop plus tokio-backed I/O providers: +// the kj::setupAsyncIo() equivalent for the tokio loop. One per thread. +TokioAsyncIoContext setupTokioAsyncIo(); + +// Resolves when the process receives signal `signum`: the tokio-loop replacement for +// kj::UnixEventPort::onSignal() (workerd's SIGTERM graceful drain). Must be awaited on the +// thread owning the TokioEventPort. Unlike UnixEventPort, KJ does not block/capture the signal +// beforehand: the tokio handler is registered when the promise is first polled, so a signal +// delivered before the event loop first runs takes its default disposition (see signal.rs). +// On Windows, SIGTERM/SIGINT are mapped to the ctrl_shutdown/ctrl_c console control events; +// the promise rejects for other signums. +kj::Promise onSignal(int signum); + +} // namespace kj_rs_io diff --git a/src/rust/cxx/kj-rs-io/error.rs b/src/rust/cxx/kj-rs-io/error.rs new file mode 100644 index 00000000000..1742cf2085b --- /dev/null +++ b/src/rust/cxx/kj-rs-io/error.rs @@ -0,0 +1,130 @@ +//! Error mapping: `std::io::Error` -> `kj::Exception`, preserving KJ's exception-type taxonomy. +//! +//! kj-http and capnp RPC change behavior based on `kj::Exception::Type` (e.g. `DISCONNECTED` +//! failures are treated as clean peer hangups rather than bugs), so the mapping of connection +//! errors matters for behavioral parity with `kj::setupAsyncIo()`. + +use cxx::IntoKjException; +use cxx::KjError; +use cxx::KjException; +use cxx::KjExceptionType; + +pub type Result = std::result::Result; + +/// An `std::io::Error` (plus operation context) that converts into a `kj::Exception` with an +/// appropriate exception type. +#[derive(Debug)] +pub struct KjIoError { + /// Name of the failing operation, included in the exception description the way KJ's + /// `KJ_SYSCALL` includes the syscall name (e.g. "`connect()`: Connection refused ..."). + op: &'static str, + inner: std::io::Error, +} + +impl KjIoError { + pub(crate) fn other(op: &'static str, message: impl std::fmt::Display) -> Self { + Self { + op, + inner: std::io::Error::other(message.to_string()), + } + } +} + +/// Attaches an operation name to `io::Error`s, for use with `Result::map_err`. +pub fn op(name: &'static str) -> impl Fn(std::io::Error) -> KjIoError { + move |inner| KjIoError { op: name, inner } +} + +fn exception_type(error: &std::io::Error) -> KjExceptionType { + use std::io::ErrorKind; + // Primary classification: by raw errno, mirroring KJ's own table (`typeOfErrno()` in + // kj/debug.c++) errno-for-errno. Consumers (kj-http, capnp-rpc) change behavior on the + // exception type, so the classes must match `kj::setupAsyncIo()` exactly — e.g. ETIMEDOUT + // is OVERLOADED in KJ (retry-later), NOT DISCONNECTED (clean peer hangup), and std's + // `ErrorKind` buckets have no stable kinds at all for KJ's fd/memory-exhaustion OVERLOADED + // set (EMFILE/ENFILE/ENOBUFS/...), hence the raw match. + #[cfg(unix)] + if let Some(errno) = error.raw_os_error() { + return errno_exception_type(errno); + } + // Fallback for synthetic (non-OS) errors — and all errors on Windows, where + // `raw_os_error()` is a Win32/WSA code, not an errno (KJ classifies those in + // `typeOfWin32Error()`; the buckets below agree with it for the kinds tokio surfaces). + match error.kind() { + // KJ's DISCONNECTED class: connection teardown, treated as a clean peer hangup. + ErrorKind::ConnectionRefused + | ErrorKind::ConnectionReset + | ErrorKind::ConnectionAborted + | ErrorKind::BrokenPipe + | ErrorKind::NotConnected + | ErrorKind::UnexpectedEof + | ErrorKind::HostUnreachable + | ErrorKind::NetworkUnreachable + | ErrorKind::NetworkDown => KjExceptionType::Disconnected, + // KJ's OVERLOADED class: temporary lack of resources (ETIMEDOUT/WSAETIMEDOUT and + // ENOMEM land here in KJ's tables). + ErrorKind::TimedOut | ErrorKind::OutOfMemory => KjExceptionType::Overloaded, + ErrorKind::Unsupported => KjExceptionType::Unimplemented, + _ => KjExceptionType::Failed, + } +} + +/// Exact mirror of KJ's `typeOfErrno()` (kj/debug.c++), so `kj::Exception::Type` matches the +/// native `kj::setupAsyncIo()` backend errno-for-errno. +#[cfg(unix)] +fn errno_exception_type(errno: i32) -> KjExceptionType { + // Errnos that are `#ifdef`-conditional in KJ's table for platform reasons, mirrored here + // with `cfg`: ENONET exists only on Linux; EOPNOTSUPP aliases ENOTSUP on Linux (KJ compiles + // its case only `#if EOPNOTSUPP != ENOTSUP` — an or-pattern with both would be an + // unreachable pattern there). + #[cfg(any(target_os = "linux", target_os = "android"))] + if errno == libc::ENONET { + return KjExceptionType::Disconnected; + } + #[cfg(not(any(target_os = "linux", target_os = "android")))] + if errno == libc::EOPNOTSUPP { + return KjExceptionType::Unimplemented; + } + match errno { + // OVERLOADED: the call failed because of a temporary lack of resources. + libc::EDQUOT + | libc::EMFILE + | libc::ENFILE + | libc::ENOBUFS + | libc::ENOLCK + | libc::ENOMEM + | libc::ENOSPC + | libc::ETIMEDOUT + | libc::EUSERS => KjExceptionType::Overloaded, + // DISCONNECTED: communication over a connection that has been lost. + libc::ENOTCONN + | libc::ECONNABORTED + | libc::ECONNREFUSED + | libc::ECONNRESET + | libc::EHOSTDOWN + | libc::EHOSTUNREACH + | libc::ENETDOWN + | libc::ENETRESET + | libc::ENETUNREACH + | libc::EPIPE => KjExceptionType::Disconnected, + // UNIMPLEMENTED: the "not supported" family (ENOTSOCK is really "syscall not + // implemented for non-sockets", per KJ's own comment). + libc::ENOSYS | libc::ENOTSUP | libc::ENOPROTOOPT | libc::ENOTSOCK => { + KjExceptionType::Unimplemented + } + _ => KjExceptionType::Failed, + } +} + +impl From for KjError { + fn from(error: KjIoError) -> Self { + let description = format!("{}: {}", error.op, error.inner); + Self::new(exception_type(&error.inner), description) + } +} + +impl IntoKjException for KjIoError { + fn into_kj_exception(self, file: &str, line: u32) -> KjException { + KjError::from(self).into_kj_exception(file, line) + } +} diff --git a/src/rust/cxx/kj-rs-io/ffi.rs b/src/rust/cxx/kj-rs-io/ffi.rs new file mode 100644 index 00000000000..f247dbb18a7 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/ffi.rs @@ -0,0 +1,832 @@ +//! The FFI island of kj-rs-io: the `#[cxx::bridge]` wire plus the crate's hand-written `unsafe` +//! boundary, in one dedicated file (file-top `#![allow(unsafe_code)]`). +//! +//! It holds: +//! +//! - the `#[cxx::bridge] mod bridge` (namespace `kj_rs_io`) — the cxx-generated C++ <-> Rust wire, +//! re-exported as `crate::ffi::*`; and +//! - every hand-written `unsafe` the macro does not generate, so the serve / net / stream modules +//! can carry the crate-root `#![deny(unsafe_code)]` and be *compiler-proven* free of hand-written +//! unsafe. Three kinds of raw thing cross the FFI boundary and are laundered into safe Rust types +//! here: +//! +//! 1. **Raw OS socket handles** arriving from C++ as an `i64` (a Unix fd or a win32 `SOCKET`; +//! see [`own_socket_from_raw`], the one platform conversion point) — the C++ side has +//! already normalized KJ's `TAKE_OWNERSHIP` / `ALREADY_CLOEXEC` flags, dup'ing when not +//! transferring ownership — plus [`dup_raw_fd`] (dup a borrowed fd into an owned one; the +//! unix-only handle tier of [`take_kj_socket`]) and [`own_fd_from_raw`] (the unix-only +//! pipe tier's `i32` fds). +//! 2. **`struct sockaddr` bytes** crossing in both directions — [`sockaddr_to_bytes`] / +//! [`sockaddr_from_bytes`]. +//! 3. **Raw `getsockopt(2)` / `setsockopt(2)`** — the socket-option passthrough behind +//! `kj::AsyncIoStream` / `kj::ConnectionReceiver` carries caller-owned option buffers whose +//! length semantics (socklen in/out) no safe std/socket2 API expresses, so the raw syscalls +//! live here ([`stream_getsockopt`] and friends). +//! +//! Because this file opts back into `unsafe` (`#![allow(unsafe_code)]`), it is the one module in +//! the crate whose soundness must be audited by hand; the rest is enforced by the compiler. +#![allow(unsafe_code)] + +use core::pin::Pin; + +/// Opaque binding of `kj::AsyncIoStream`, and the cxx-bridged operations on it (shared-receiver +/// shims; safe to call). Re-exported as `crate::ffi::*` so the rest of the crate (and the +/// crate-root re-exports) keep using `ffi::`. +pub use bridge::KjAsyncIoStream; +pub use bridge::kj_stream_get_handle; +pub use bridge::unwrap_tokio_stream; +use cxx::KjException; +use kj_rs::KjOwn; + +use crate::error::Result; +use crate::error::op; +use crate::net::TokioAddress; +use crate::net::TokioListener; +use crate::net::address_clone; +use crate::net::address_connect_index; +use crate::net::address_count; +use crate::net::address_listen; +use crate::net::address_raw_sockaddr; +use crate::net::address_to_string; +use crate::net::listener_accept; +use crate::net::listener_local_addr; +use crate::net::listener_port; +use crate::net::network_get_sockaddr; +use crate::net::network_parse_address; +use crate::net::wrap_connecting_socket_fd; +use crate::net::wrap_listen_fd; +use crate::net::wrap_socket_fd; +use crate::readiness::wait_fd_readable; +use crate::signal::wait_for_signal; +use crate::stream::TokioInputFd; +use crate::stream::TokioOutputFd; +use crate::stream::TokioStream; +use crate::stream::input_fd_try_read; +use crate::stream::output_fd_write; +use crate::stream::stream_local_addr; +use crate::stream::stream_peer_addr; +use crate::stream::stream_raw_handle; +use crate::stream::stream_shutdown_write; +use crate::stream::stream_take; +use crate::stream::stream_try_read; +use crate::stream::stream_when_write_disconnected; +use crate::stream::stream_write; +use crate::stream::wrap_input_fd; +use crate::stream::wrap_output_fd; + +#[cxx::bridge(namespace = "kj_rs_io")] +// FFI island: the cxx bridge macro generates the `unsafe` extern shims, and this module declares +// the hand-written `unsafe extern "C++"` / `async unsafe fn` bridge surface. +// unnecessary_box_returns: returning an opaque Rust type to C++ as `Box` is the cxx idiom. +#[expect(clippy::unnecessary_box_returns)] +// missing_safety_doc fires (or not) deep inside the macro expansion depending on which bridge +// items are publicly re-exported, so an `#[expect]` could go unfulfilled. +#[expect(clippy::allow_attributes)] +#[allow(clippy::missing_safety_doc)] +mod bridge { + extern "Rust" { + type TokioStream; + type TokioListener; + type TokioAddress; + type TokioInputFd; + type TokioOutputFd; + + // ================================================================================== + // Streams (TCP or Unix domain, behind kj::AsyncIoStream) + + /// Reads until at least `min_bytes` are available (or EOF), up to `buf.len()`. Returns + /// the number of bytes read; fewer than `min_bytes` indicates EOF. The kj `tryRead` + /// contract. + async unsafe fn stream_try_read<'a>( + stream: &'a TokioStream, + buf: &'a mut [u8], + min_bytes: usize, + ) -> Result; + + /// Writes the entire buffer (write-all semantics). + async unsafe fn stream_write<'a>(stream: &'a TokioStream, buf: &'a [u8]) -> Result<()>; + + /// Resolves when the stream has become disconnected such that new writes will fail. + /// See `TokioStream::when_write_disconnected` for the mechanism and platform caveats. + async unsafe fn stream_when_write_disconnected<'a>(stream: &'a TokioStream) -> Result<()>; + + /// `shutdown(SHUT_WR)`: cleanly shut down the write end, keeping the read end open. + fn stream_shutdown_write(stream: &TokioStream) -> Result<()>; + + /// The underlying raw OS socket handle (fd on unix, `SOCKET` on windows) as an `i64`, + /// backing `kj::AsyncIoStream::getFd()` / `getWin32Handle()`. + fn stream_raw_handle(stream: &TokioStream) -> Result; + + /// Raw `struct sockaddr` bytes of the socket's locally-bound address (the + /// `getsockname()` passthrough). + fn stream_local_addr(stream: &TokioStream) -> Result>; + + /// Raw `struct sockaddr` bytes of the connected peer's address (the `getpeername()` + /// passthrough, also used by the accept-loop peer-filter check). + fn stream_peer_addr(stream: &TokioStream) -> Result>; + + /// `getsockopt(2)` on the underlying socket. `value.len()` is the caller's in-length + /// (the kernel truncates the option value to it); returns the syscall's reported + /// out-length, which the caller must mirror back exactly (raw socklen in/out + /// semantics). + fn stream_getsockopt( + stream: &TokioStream, + level: i32, + option: i32, + value: &mut [u8], + ) -> Result; + + /// `setsockopt(2)` on the underlying socket. + fn stream_setsockopt( + stream: &TokioStream, + level: i32, + option: i32, + value: &[u8], + ) -> Result<()>; + + /// Moves the native stream out, leaving `stream` hollow (all further ops error). + /// Unsafe contract: no I/O futures may currently borrow `stream`. + fn stream_take(stream: &mut TokioStream) -> Result>; + + // ================================================================================== + // Network addresses (kj::Network::parseAddress grammar subset) + + /// Parses a KJ address string ("1.2.3.4:80", "[::1]:80", "host:80", "*", "*:80", + /// "unix:/path"), resolving hostnames via DNS. `port_hint` fills in a missing port. + /// (Owned `String`: the future must not borrow the caller's buffer across the DNS + /// suspension.) + async fn network_parse_address(addr: String, port_hint: u16) -> Result>; + + /// Builds an address from a raw `struct sockaddr` (AF_INET / AF_INET6 / AF_UNIX). + fn network_get_sockaddr(sockaddr: &[u8]) -> Result>; + + /// Connects to exactly the `index`th resolved address (no fallback). The C++ side + /// drives the try-each-address loop so it can apply restrictPeers() filtering per + /// address (KJ parity). + async unsafe fn address_connect_index<'a>( + addr: &'a TokioAddress, + index: usize, + ) -> Result>; + + /// Number of resolved socket addresses behind this address (>= 1). + fn address_count(addr: &TokioAddress) -> usize; + + /// Raw `struct sockaddr` bytes of the `index`th resolved address, for C++-side + /// kj::_::NetworkFilter (restrictPeers) checks. + fn address_raw_sockaddr(addr: &TokioAddress, index: usize) -> Result>; + + /// Binds + listens on the (first) address. Wildcard addresses bind dual-stack. + fn address_listen(addr: &TokioAddress) -> Result>; + + fn address_clone(addr: &TokioAddress) -> Box; + fn address_to_string(addr: &TokioAddress) -> String; + + // ================================================================================== + // Listeners (kj::ConnectionReceiver) + + async unsafe fn listener_accept<'a>( + listener: &'a TokioListener, + ) -> Result>; + + /// The locally-bound port (0 for Unix domain sockets, mirroring KJ). + fn listener_port(listener: &TokioListener) -> Result; + + /// Raw `struct sockaddr` bytes of the listener's bound address (the `getsockname()` + /// passthrough). + fn listener_local_addr(listener: &TokioListener) -> Result>; + + /// `getsockopt(2)` on the listening socket; same length semantics as + /// `stream_getsockopt`. + fn listener_getsockopt( + listener: &TokioListener, + level: i32, + option: i32, + value: &mut [u8], + ) -> Result; + + /// `setsockopt(2)` on the listening socket. + fn listener_setsockopt( + listener: &TokioListener, + level: i32, + option: i32, + value: &[u8], + ) -> Result<()>; + + // ================================================================================== + // Socket-handle wrapping (kj::LowLevelAsyncIoProvider). + // + // The `i64` is a raw OS socket handle: a Unix fd or a win32 `SOCKET` (`i64` fits both + // losslessly, with `-1` ≡ `INVALID_SOCKET` as the shared sentinel). All of these take + // ownership of the handle. The C++ side normalizes + // TAKE_OWNERSHIP/ALREADY_CLOEXEC/ALREADY_NONBLOCK flags (dup'ing when not taking + // ownership) before calling in; [`own_socket_from_raw`] is the single point where the + // raw handle becomes an owned socket. + + /// Wraps a connected stream socket handle (TCP or Unix domain, detected automatically). + fn wrap_socket_fd(handle: i64) -> Result>; + + /// Wraps a bound+listening socket handle (TCP or Unix domain, detected automatically). + fn wrap_listen_fd(handle: i64) -> Result>; + + /// Wraps an unconnected TCP socket handle and connects it to `sockaddr` (a raw + /// `struct sockaddr`, AF_INET/AF_INET6 only; owned copy, since the caller's pointer + /// need not outlive the call). + async fn wrap_connecting_socket_fd( + handle: i64, + sockaddr: Vec, + ) -> Result>; + + /// Wraps a readable fd (pipe, character device, socket). Regular files are rejected by + /// the OS readiness API (same as KJ's epoll-based provider). Unix only (the pipe tier + /// keeps `i32` fds). + fn wrap_input_fd(fd: i32) -> Result>; + + async unsafe fn input_fd_try_read<'a>( + stream: &'a TokioInputFd, + buf: &'a mut [u8], + min_bytes: usize, + ) -> Result; + + /// Wraps a writable fd (pipe, character device, socket). + fn wrap_output_fd(fd: i32) -> Result>; + + async unsafe fn output_fd_write<'a>(stream: &'a TokioOutputFd, buf: &'a [u8]) + -> Result<()>; + + // ================================================================================== + // Signals (kj::UnixEventPort::onSignal replacement; see signal.rs for semantics) + + /// Resolves when the process receives signal `signum` (on Windows: the mapped + /// SIGTERM/SIGINT console control event). + async fn wait_for_signal(signum: i32) -> Result<()>; + + // ================================================================================== + // Fd readiness (kj::UnixEventPort::FdObserver::whenBecomesReadable replacement, + // backing kj_rs_io::FileWatcher in file-watcher.h; see readiness.rs for semantics) + + /// Resolves when `fd` becomes readable (readiness already pending at call time is + /// reported immediately). The caller must keep `fd` open until the returned promise + /// resolves or is dropped, and must not watch the same fd twice concurrently. + /// Unix only. + async fn wait_fd_readable(fd: i32) -> Result<()>; + } + + unsafe extern "C++" { + include!("kj-rs-io/unwrap.h"); + + /// `kj::AsyncIoStream`, opaque. Used by [`unwrap_kj_stream`]. + #[namespace = "kj"] + #[cxx_name = "AsyncIoStream"] + type KjAsyncIoStream; + + /// Implemented in `async-io.c++`: downcasts to the kj-rs-io wrapper and moves the native + /// stream out. Throws (surfaced as `Err`) for foreign streams. + #[cxx_name = "unwrapTokioStream"] + fn unwrap_tokio_stream(stream: Pin<&mut KjAsyncIoStream>) -> Result>; + + // Bridged operations on a foreign `kj::AsyncIoStream`, backing `serve_kj_stream`'s + // duplex-pump fallback (serve.rs). Shared receivers (`&KjAsyncIoStream`, const_cast + // shims in unwrap.h): a kj two-way stream supports one concurrent read and one write, + // which the pump models as concurrent shared borrows of the stream it owns. All + // returned futures must be polled on the KJ event-loop thread owning the stream. + + /// Corresponds to `kj::AsyncIoStream::tryRead(buffer, min_bytes, buffer.len())`. + #[cxx_name = "kjStreamTryRead"] + async fn kj_stream_try_read( + stream: &KjAsyncIoStream, + buffer: &mut [u8], + min_bytes: usize, + ) -> Result; + + /// Corresponds to `kj::AsyncIoStream::write(buffer)` (write-all semantics). + #[cxx_name = "kjStreamWrite"] + async fn kj_stream_write(stream: &KjAsyncIoStream, buffer: &[u8]) -> Result<()>; + + /// Corresponds to `kj::AsyncIoStream::shutdownWrite()`. + #[cxx_name = "kjStreamShutdownWrite"] + fn kj_stream_shutdown_write(stream: &KjAsyncIoStream); + + /// The stream's underlying raw OS socket handle (fd on unix, `SOCKET` on windows; + /// `kj::AsyncIoStream::getFd()` / `getWin32Handle()`) as an `i64`, or -1 if it exposes + /// none. Backs the handle tier of [`take_kj_socket`]. + #[cxx_name = "kjStreamGetHandle"] + fn kj_stream_get_handle(stream: &KjAsyncIoStream) -> i64; + } +} + +// ====================================================================================== +// Raw socket handles / file descriptors. + +/// Materializes an owned file descriptor from a raw `i32` that arrived across the FFI bridge +/// (the unix-only pipe tier — `wrap_input_fd`/`wrap_output_fd`; sockets go through +/// [`own_socket_from_raw`]). +/// +/// Callable from safe code: the invariant it relies on (`fd` is open and its ownership has been +/// transferred to us) is structurally upheld by the cxx bridge — the C++ side normalizes KJ's +/// `TAKE_OWNERSHIP` / `ALREADY_CLOEXEC` flags and dup's the fd when the caller is not handing +/// over ownership. The returned `OwnedFd` becomes the sole owner and closes it on drop. +#[cfg(unix)] +#[must_use] +pub fn own_fd_from_raw(fd: i32) -> std::os::fd::OwnedFd { + use std::os::fd::FromRawFd; + // `OwnedFd`'s invariant is "an open fd, never -1" (-1 is its niche), so a negative value + // here would be library-level UB rather than an error. C++ callers normalize through + // prepareFd, but its all-flags-set path (TAKE_OWNERSHIP|ALREADY_CLOEXEC|ALREADY_NONBLOCK) + // performs no syscall that would catch a bad fd — enforce the contract at THE conversion + // point instead of inheriting the UB. + assert!(fd >= 0, "invalid fd crossed the FFI bridge: {fd}"); + // Safety: per the bridge contract `fd` is open and owned by us from this point on. + unsafe { std::os::fd::OwnedFd::from_raw_fd(fd) } +} + +/// Materializes an owned socket from a raw OS socket handle that arrived across the FFI bridge +/// as an `i64`: a Unix fd here, a win32 `SOCKET` in the `cfg(windows)` twin below. This is THE +/// one platform conversion point — behind it everything is a uniform `socket2::Socket` / +/// std/tokio socket type. +/// +/// Callable from safe code: the invariant it relies on (`handle` is an open socket whose +/// ownership has been transferred to us) is structurally upheld by the cxx bridge — the C++ +/// side normalizes KJ's fd-wrapping flags, dup'ing when the caller is not handing over +/// ownership. The returned socket becomes the sole owner and closes it on drop. +#[cfg(unix)] +#[must_use] +pub fn own_socket_from_raw(handle: i64) -> socket2::Socket { + // A unix fd is a non-negative int: the bridge widened it losslessly to i64, so the + // narrowing back to i32 cannot truncate for any legitimate handle. Enforce that (rejecting + // -1/garbage) here at THE conversion point rather than inheriting `OwnedFd`'s niche UB. + assert!( + (0..=i64::from(i32::MAX)).contains(&handle), + "invalid socket fd crossed the FFI bridge: {handle}" + ); + #[expect(clippy::cast_possible_truncation)] + let fd = handle as i32; + socket2::Socket::from(own_fd_from_raw(fd)) +} + +/// The `cfg(windows)` twin of [`own_socket_from_raw`]: the raw handle is a winsock `SOCKET` +/// (`u64`-shaped `RawSocket`; a live SOCKET fits in an `i64` without colliding with the -1 +/// sentinel, which is `INVALID_SOCKET` and never crosses the bridge as an owned handle). +// Validated by Windows CI; mirrors the unix arm. +#[cfg(windows)] +#[must_use] +pub fn own_socket_from_raw(handle: i64) -> socket2::Socket { + use std::os::windows::io::FromRawSocket; + use std::os::windows::io::OwnedSocket; + use std::os::windows::io::RawSocket; + // Live SOCKET values are non-negative in i64 (the -1 sentinel is INVALID_SOCKET, which + // `OwnedSocket` forbids as its niche and which must never cross the bridge as an owned + // handle) — enforce at THE conversion point rather than inheriting the niche UB. + assert!( + handle >= 0, + "invalid SOCKET crossed the FFI bridge: {handle}" + ); + // The bridge carries the SOCKET's bits verbatim. + #[allow(clippy::cast_sign_loss)] + let raw = handle as RawSocket; + // Safety: per the bridge contract `handle` is an open SOCKET owned by us from this point on. + let owned = unsafe { OwnedSocket::from_raw_socket(raw) }; + socket2::Socket::from(owned) +} + +/// Duplicates a *borrowed* raw fd into an independently-owned fd (`F_DUPFD_CLOEXEC`), for the +/// handle tier of [`take_kj_socket`]: the kj stream keeps its own fd, we get a fresh dup. +/// +/// Unix only, deliberately: no windows twin is needed. The handle tier only fires for *foreign* +/// handle-backed kj streams, and under the all-rust mode on Windows every socket-backed stream +/// originates in kj-rs-io (tier-1 unwrap); the other in-process dup users don't dup on windows +/// either (`when_write_disconnected` is unix-only — never-resolving on windows, KJ parity — and +/// windows `shutdown_write` borrows via `SockRef` instead of dup'ing). If a windows twin is +/// ever needed, `std::os::windows::io::BorrowedSocket::try_clone_to_owned` +/// (`WSADuplicateSocketW`) is the same-process equivalent. +/// +/// # Errors +/// +/// Returns the `dup()` `io::Error` (mapped to a `kj::Exception`) if the syscall fails. +#[cfg(unix)] +pub fn dup_raw_fd(fd: i32) -> Result { + // Safety: the caller (take_kj_socket) holds the kj stream, which keeps `fd` open for the + // duration of the call; we immediately dup it into an independently-owned fd. + unsafe { std::os::fd::BorrowedFd::borrow_raw(fd) } + .try_clone_to_owned() + .map_err(op("dup()")) +} + +// ====================================================================================== +// `struct sockaddr` <-> bytes. + +/// Copies a `socket2::SockAddr`'s initialized `struct sockaddr` bytes into an owned `Vec`, to +/// hand across the bridge for the C++ side's `kj::_::NetworkFilter` (restrictPeers) checks. +#[must_use] +pub fn sockaddr_to_bytes(sockaddr: &socket2::SockAddr) -> Vec { + // Safety: as_ptr()/len() delimit an initialized sockaddr owned by `sockaddr`. + let bytes = unsafe { + std::slice::from_raw_parts(sockaddr.as_ptr().cast::(), sockaddr.len() as usize) + }; + bytes.to_vec() +} + +/// Decodes raw `struct sockaddr` bytes (arriving from C++) into a `socket2::SockAddr`. +/// +/// # Errors +/// +/// Errors if the byte length is zero or exceeds `sockaddr_storage`. +#[cfg(any(unix, windows))] +pub fn sockaddr_from_bytes(bytes: &[u8]) -> Result { + use crate::error::KjIoError; + + let mut storage = socket2::SockAddrStorage::zeroed(); + let storage_size = std::mem::size_of::(); + if bytes.is_empty() || bytes.len() > storage_size { + return Err(KjIoError::other("sockaddr", "invalid sockaddr length")); + } + // Safety: SockAddrStorage is plain-old-data large enough for any sockaddr; we copy + // `bytes.len() <= size_of::()` bytes into it. + unsafe { + std::ptr::copy_nonoverlapping( + bytes.as_ptr(), + std::ptr::from_mut(&mut storage).cast::(), + bytes.len(), + ); + } + #[expect(clippy::cast_possible_truncation)] + let len = bytes.len() as socket2::socklen_t; + // Safety: `storage` is a zeroed sockaddr_storage with the caller's `len` bytes copied in, + // satisfying SockAddr::new's layout/length requirements. The address *family* is NOT + // validated here: for families socket2 does not understand (e.g. AF_NETLINK), accessors + // like `as_socket()`/`as_pathname()` return `None`, and callers must surface that as an + // "unsupported sockaddr family" error (see `net.rs::network_get_sockaddr`) rather than + // assume a known family. + Ok(unsafe { socket2::SockAddr::new(storage, len) }) +} + +// ====================================================================================== +// Raw `getsockopt(2)` / `setsockopt(2)`. +// +// The socket-option passthrough behind `kj::AsyncIoStream::get/setsockopt` and +// `kj::ConnectionReceiver::get/setsockopt`. The option buffer is caller-owned opaque bytes with +// raw socklen in/out semantics (the caller's buffer may be smaller than the option value, and the +// syscall's reported length must be surfaced verbatim), which no safe std/socket2 API expresses — +// so the raw syscalls are declared and called here, in the unsafe island. + +/// Raw `getsockopt(2)` on a borrowed socket fd. `value.len()` is passed as the in `optlen` (the +/// kernel truncates the option value to it); the syscall's reported out `optlen` is returned so +/// the C++ caller can mirror `*length = socklen` exactly as `KJ_SYSCALL(::getsockopt(...))` did. +#[cfg(unix)] +fn getsockopt_raw( + fd: std::os::fd::BorrowedFd<'_>, + level: i32, + option: i32, + value: &mut [u8], +) -> Result { + use core::ffi::c_int; + use core::ffi::c_void; + use std::os::fd::AsRawFd; + unsafe extern "C" { + fn getsockopt( + sockfd: c_int, + level: c_int, + optname: c_int, + optval: *mut c_void, + optlen: *mut socket2::socklen_t, + ) -> c_int; + } + #[expect(clippy::cast_possible_truncation)] + let mut optlen = value.len() as socket2::socklen_t; + // Safety: simple syscall wrapper. `fd` is a live socket fd (borrowed from the tokio object + // for the duration of the call); `value.as_mut_ptr()` with in-`optlen == value.len()` + // delimits writable caller memory the kernel fills (never past `optlen`); `&raw mut optlen` + // is a valid in/out pointer for the call. + let rc = unsafe { + getsockopt( + fd.as_raw_fd(), + level, + option, + value.as_mut_ptr().cast::(), + &raw mut optlen, + ) + }; + if rc != 0 { + return Err(op("getsockopt()")(std::io::Error::last_os_error())); + } + Ok(optlen as usize) +} + +/// Raw `setsockopt(2)` on a borrowed socket fd. +#[cfg(unix)] +fn setsockopt_raw( + fd: std::os::fd::BorrowedFd<'_>, + level: i32, + option: i32, + value: &[u8], +) -> Result<()> { + use core::ffi::c_int; + use core::ffi::c_void; + use std::os::fd::AsRawFd; + unsafe extern "C" { + fn setsockopt( + sockfd: c_int, + level: c_int, + optname: c_int, + optval: *const c_void, + optlen: socket2::socklen_t, + ) -> c_int; + } + #[expect(clippy::cast_possible_truncation)] + let optlen = value.len() as socket2::socklen_t; + // Safety: simple syscall wrapper. `fd` is a live socket fd (borrowed from the tokio object + // for the duration of the call); `value.as_ptr()` with `optlen == value.len()` delimits + // readable caller memory the kernel only reads. + let rc = unsafe { + setsockopt( + fd.as_raw_fd(), + level, + option, + value.as_ptr().cast::(), + optlen, + ) + }; + if rc != 0 { + return Err(op("setsockopt()")(std::io::Error::last_os_error())); + } + Ok(()) +} + +/// Raw ws2_32 `getsockopt` on a borrowed `SOCKET`. Same socklen in/out semantics as the unix +/// arm above: `value.len()` is passed as the in `optlen`, and the reported out `optlen` is +/// returned verbatim. +// Validated by Windows CI; mirrors the unix arm. +#[cfg(windows)] +fn getsockopt_raw( + sock: std::os::windows::io::BorrowedSocket<'_>, + level: i32, + option: i32, + value: &mut [u8], +) -> Result { + use core::ffi::c_char; + use core::ffi::c_int; + use std::os::windows::io::AsRawSocket; + use std::os::windows::io::RawSocket; + #[link(name = "ws2_32")] + unsafe extern "system" { + fn getsockopt( + s: RawSocket, + level: c_int, + optname: c_int, + optval: *mut c_char, + optlen: *mut c_int, + ) -> c_int; + } + #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] + let mut optlen = value.len() as c_int; + // Safety: simple syscall wrapper. `sock` is a live `SOCKET` (borrowed from the tokio object + // for the duration of the call); `value.as_mut_ptr()` with in-`optlen == value.len()` + // delimits writable caller memory winsock fills (never past `optlen`); `&raw mut optlen` is + // a valid in/out pointer for the call. + let rc = unsafe { + getsockopt( + sock.as_raw_socket(), + level, + option, + value.as_mut_ptr().cast::(), + &raw mut optlen, + ) + }; + if rc != 0 { + // rc is SOCKET_ERROR (-1); `last_os_error()` reads `WSAGetLastError()` on Windows. + return Err(op("getsockopt()")(std::io::Error::last_os_error())); + } + // The out-length winsock reports is non-negative (and bounded by the in-length). + #[allow(clippy::cast_sign_loss)] + let reported = optlen as usize; + Ok(reported) +} + +/// Raw ws2_32 `setsockopt` on a borrowed `SOCKET`. +// Validated by Windows CI; mirrors the unix arm. +#[cfg(windows)] +fn setsockopt_raw( + sock: std::os::windows::io::BorrowedSocket<'_>, + level: i32, + option: i32, + value: &[u8], +) -> Result<()> { + use core::ffi::c_char; + use core::ffi::c_int; + use std::os::windows::io::AsRawSocket; + use std::os::windows::io::RawSocket; + #[link(name = "ws2_32")] + unsafe extern "system" { + fn setsockopt( + s: RawSocket, + level: c_int, + optname: c_int, + optval: *const c_char, + optlen: c_int, + ) -> c_int; + } + #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)] + let optlen = value.len() as c_int; + // Safety: simple syscall wrapper. `sock` is a live `SOCKET` (borrowed from the tokio object + // for the duration of the call); `value.as_ptr()` with `optlen == value.len()` delimits + // readable caller memory winsock only reads. + let rc = unsafe { + setsockopt( + sock.as_raw_socket(), + level, + option, + value.as_ptr().cast::(), + optlen, + ) + }; + if rc != 0 { + // rc is SOCKET_ERROR (-1); `last_os_error()` reads `WSAGetLastError()` on Windows. + return Err(op("setsockopt()")(std::io::Error::last_os_error())); + } + Ok(()) +} + +pub fn stream_getsockopt( + stream: &TokioStream, + level: i32, + option: i32, + value: &mut [u8], +) -> Result { + #[cfg(unix)] + { + getsockopt_raw(stream.as_borrowed_fd()?, level, option, value) + } + // Validated by Windows CI; mirrors the unix arm. + #[cfg(windows)] + { + getsockopt_raw(stream.as_borrowed_socket()?, level, option, value) + } + #[cfg(not(any(unix, windows)))] + { + let _ = (stream, level, option, value); + Err(crate::error::KjIoError::other( + "getsockopt", + "not implemented by kj-rs-io on this platform", + )) + } +} + +pub fn stream_setsockopt( + stream: &TokioStream, + level: i32, + option: i32, + value: &[u8], +) -> Result<()> { + #[cfg(unix)] + { + setsockopt_raw(stream.as_borrowed_fd()?, level, option, value) + } + // Validated by Windows CI; mirrors the unix arm. + #[cfg(windows)] + { + setsockopt_raw(stream.as_borrowed_socket()?, level, option, value) + } + #[cfg(not(any(unix, windows)))] + { + let _ = (stream, level, option, value); + Err(crate::error::KjIoError::other( + "setsockopt", + "not implemented by kj-rs-io on this platform", + )) + } +} + +pub fn listener_getsockopt( + listener: &TokioListener, + level: i32, + option: i32, + value: &mut [u8], +) -> Result { + #[cfg(unix)] + { + getsockopt_raw(listener.as_borrowed_fd(), level, option, value) + } + // Validated by Windows CI; mirrors the unix arm. + #[cfg(windows)] + { + getsockopt_raw(listener.as_borrowed_socket(), level, option, value) + } + #[cfg(not(any(unix, windows)))] + { + let _ = (listener, level, option, value); + Err(crate::error::KjIoError::other( + "getsockopt", + "not implemented by kj-rs-io on this platform", + )) + } +} + +pub fn listener_setsockopt( + listener: &TokioListener, + level: i32, + option: i32, + value: &[u8], +) -> Result<()> { + #[cfg(unix)] + { + setsockopt_raw(listener.as_borrowed_fd(), level, option, value) + } + // Validated by Windows CI; mirrors the unix arm. + #[cfg(windows)] + { + setsockopt_raw(listener.as_borrowed_socket(), level, option, value) + } + #[cfg(not(any(unix, windows)))] + { + let _ = (listener, level, option, value); + Err(crate::error::KjIoError::other( + "setsockopt", + "not implemented by kj-rs-io on this platform", + )) + } +} + +// ====================================================================================== +// Typed read/write halves of a pumped `kj::AsyncIoStream`. +// +// kj's stream contract — at most one read and one write may be in flight at once — is prose in +// kj; these halves make the borrow checker enforce it. Each half's operations take `&mut self`, +// so an in-flight operation's future exclusively borrows its half (a second overlapping read is +// a compile error), and `split_kj_stream` takes the owner's `&mut`, so while the halves live +// nothing else (an unwrap, another split) can touch the stream. The bridged operations behind +// them are not re-exported: the halves are the only way to drive a foreign stream. + +/// The read direction of a pumped stream. See the module comment above. +pub struct KjStreamReadHalf<'a>(&'a KjAsyncIoStream); + +/// The write direction of a pumped stream (writes and the write-side shutdown). See the module +/// comment above. +pub struct KjStreamWriteHalf<'a>(&'a KjAsyncIoStream); + +/// Splits the owned stream into its two directions. Holding the owner's `&mut` for the halves' +/// lifetime proves exactly one pair exists and reserves the stream for them. +// The unused `&mut` is the point (see the doc comment): it reserves the stream for the halves. +#[expect(clippy::needless_pass_by_ref_mut)] +pub fn split_kj_stream( + stream: &mut KjOwn, +) -> (KjStreamReadHalf<'_>, KjStreamWriteHalf<'_>) { + let stream = &**stream; + (KjStreamReadHalf(stream), KjStreamWriteHalf(stream)) +} + +impl KjStreamReadHalf<'_> { + /// `kj::AsyncIoStream::tryRead(buffer, min_bytes, buffer.len())`. + // The unused `&mut self` is the point: an in-flight read's future exclusively borrows the + // read half (kj's one-read-in-flight contract), see the section comment above. + #[expect(clippy::needless_pass_by_ref_mut)] + pub(crate) async fn try_read( + &mut self, + buf: &mut [u8], + min_bytes: usize, + ) -> std::result::Result { + bridge::kj_stream_try_read(self.0, buf, min_bytes).await + } +} + +impl KjStreamWriteHalf<'_> { + /// `kj::AsyncIoStream::write(buffer)` (write-all semantics). + // The unused `&mut self` is the point: an in-flight write's future exclusively borrows the + // write half (kj's one-write-in-flight contract), see the section comment above. + #[expect(clippy::needless_pass_by_ref_mut)] + pub(crate) async fn write(&mut self, buf: &[u8]) -> std::result::Result<(), KjException> { + bridge::kj_stream_write(self.0, buf).await + } + + /// `kj::AsyncIoStream::shutdownWrite()`. + // The unused `&mut self` is the point: an exclusive borrow of the write half serializes + // write-side operations (kj's one-write-in-flight contract), see the section comment above. + #[expect(clippy::needless_pass_by_ref_mut)] + pub(crate) fn shutdown_write(&mut self) { + bridge::kj_stream_shutdown_write(self.0); + } +} + +// ====================================================================================== +// Pub `unsafe fn` FFI entry point (`Pin<&mut kj::AsyncIoStream>`). +// +// The owning native-serve entry points (`take_kj_socket`, `serve_kj_stream`) are safe fns in +// [`crate::serve`] — ownership arrives as a `KjOwn` and no raw pointer crosses the crate's +// public surface. Only the borrow-based unwrap remains here: C++ keeps the (hollow) wrapper, +// so its "no I/O in flight" precondition cannot be expressed structurally. + +/// Recovers the native [`TokioStream`] out of a `kj::AsyncIoStream` that was created by +/// kj-rs-io, leaving the C++ wrapper hollow (any further I/O through it fails). +/// +/// # Errors +/// +/// Returns an error if the stream is not a kj-rs-io tokio-backed stream, or was already +/// unwrapped. +/// +/// # Safety +/// +/// No I/O operations (reads, writes, `whenWriteDisconnected`) may be in flight on the stream: +/// their futures borrow the same native object this function moves out. +pub unsafe fn unwrap_kj_stream( + stream: Pin<&mut KjAsyncIoStream>, +) -> std::result::Result, KjException> { + bridge::unwrap_tokio_stream(stream) +} diff --git a/src/rust/cxx/kj-rs-io/file-watcher.c++ b/src/rust/cxx/kj-rs-io/file-watcher.c++ new file mode 100644 index 00000000000..c37369593fa --- /dev/null +++ b/src/rust/cxx/kj-rs-io/file-watcher.c++ @@ -0,0 +1,213 @@ +// kj_rs_io::FileWatcher implementation. Each platform backend is a line-for-line port of +// workerd's kj-mode FileWatcher (workerd.c++), with the one difference that waiting for the +// notification fd to become readable goes through tokio's AsyncFd (wait_fd_readable, see +// readiness.rs) instead of kj::UnixEventPort::FdObserver::whenBecomesReadable(). Everything +// else -- which fds get created, which watch masks are used, how events are drained and +// filtered -- is kept identical so that --watch behaves the same on either event loop. + +#include "kj-rs-io/file-watcher.h" + +#include "kj-rs-io/ffi.rs.h" + +#include +#include +#include +#include + +#include + +#if __linux__ +#include +#include +#include +#elif __APPLE__ || __FreeBSD__ || __OpenBSD__ || __NetBSD__ || __DragonFly__ +#define KJ_RS_IO_USE_KQUEUE_FOR_FILE_WATCHER 1 +#include +#include +#include +#include +#include +#endif + +namespace kj_rs_io { + +#if __linux__ + +// inotify backend. Watches each file's parent directory (IN_DELETE | IN_MODIFY | IN_MOVE | +// IN_CREATE) and filters events by basename, so files replaced by rename (editors' atomic +// saves) or deleted-and-recreated keep firing, and the watched file itself need not exist yet. +struct FileWatcher::Impl { + kj::OwnFd inotifyFd; + + kj::HashMap watches; + kj::HashMap> filesWatched; + + Impl(): inotifyFd(KJ_SYSCALL_FD(inotify_init1(IN_NONBLOCK | IN_CLOEXEC))) {} + + bool isSupported() { + return true; + } + + void watch(kj::PathPtr path, kj::Maybe file) { + // The inotify backend doesn't use `file`; it watches the parent directory. + + auto pathStr = path.parent().toNativeString(true); + + int wd = watches.findOrCreate(pathStr, [&]() { + int wd; + uint32_t mask = IN_DELETE | IN_MODIFY | IN_MOVE | IN_CREATE; + KJ_SYSCALL(wd = inotify_add_watch(inotifyFd, pathStr.cStr(), mask)); + return decltype(watches)::Entry{kj::mv(pathStr), wd}; + }); + + auto &files = + filesWatched.findOrCreate(wd, [&]() { return decltype(filesWatched)::Entry{wd, {}}; }); + + files.upsert(kj::str(path.basename()[0]), [](auto &&...) {}); + } + + kj::Promise onChange() { + kj::byte buffer[4096]{}; + + for (;;) { + ssize_t n; + KJ_NONBLOCKING_SYSCALL(n = read(inotifyFd, buffer, sizeof(buffer))); + + if (n < 0) { + // No more data to read: wait for the inotify fd to become readable again. + co_await wait_fd_readable(inotifyFd.get()); + continue; + } + + kj::byte *ptr = buffer; + while (n > 0) { + KJ_ASSERT(n >= sizeof(struct inotify_event)); + + auto &event = *reinterpret_cast(ptr); + size_t eventSize = sizeof(struct inotify_event) + event.len; + KJ_ASSERT(n >= eventSize); + KJ_ASSERT(eventSize % sizeof(void *) == 0); + ptr += eventSize; + n -= eventSize; + + if (event.len > 0 && event.name[0] != '\0') { + auto &watched = KJ_ASSERT_NONNULL(filesWatched.find(event.wd)); + if (watched.find(kj::StringPtr(event.name)) != kj::none) { + // HIT! We saw a change. + co_return; + } + } + } + } + } +}; + +#elif KJ_RS_IO_USE_KQUEUE_FOR_FILE_WATCHER + +// kqueue backend. One EVFILT_VNODE registration per watched file (dup of the already-open +// config fd when available, else opened by path -- so the path must exist). NOTE_DELETE / +// NOTE_RENAME on the old inode cover atomic-rename saves. kqueue doesn't scale to whole +// directory trees, but we only watch the specific files opened while parsing the config. +struct FileWatcher::Impl { + kj::OwnFd kqueueFd; + kj::Vector filesWatched; + + Impl(): kqueueFd(makeKqueue()) {} + + bool isSupported() { + return true; + } + + void watch(kj::PathPtr path, kj::Maybe file) { + KJ_IF_SOME(f, file) { + KJ_IF_SOME(fd, f.getFd()) { + // We need to duplicate the fd because the original will probably be closed later, and + // closing the fd unregisters it from kqueue. + watchFd(KJ_SYSCALL_FD(dup(fd))); + return; + } + } + + // No existing file, open from disk. + watchFd(KJ_SYSCALL_FD(open(path.toNativeString(true).cStr(), O_RDONLY))); + } + + kj::Promise onChange() { + for (;;) { + struct kevent event; + struct timespec timeout; + memset(&event, 0, sizeof(event)); + memset(&timeout, 0, sizeof(timeout)); + + int n; + KJ_SYSCALL(n = kevent(kqueueFd, nullptr, 0, &event, 1, &timeout)); + + if (n == 0) { + // No events: wait for the kqueue fd to become readable, indicating an event has been + // delivered. + co_await wait_fd_readable(kqueueFd.get()); + continue; + } else { + // We only registered for events that indicate changes in the first place, so there's + // no need to examine the event: it definitely means something changed. + co_return; + } + } + } + + static kj::OwnFd makeKqueue() { + auto fd = KJ_SYSCALL_FD(kqueue()); + KJ_SYSCALL(fcntl(fd, F_SETFD, FD_CLOEXEC)); + return kj::mv(fd); + } + + void watchFd(kj::OwnFd fd) { + KJ_SYSCALL(fcntl(fd, F_SETFD, FD_CLOEXEC)); + + struct kevent change; + memset(&change, 0, sizeof(change)); + change.ident = fd.get(); + change.filter = EVFILT_VNODE; + change.flags = EV_ADD | EV_CLEAR; + change.fflags = NOTE_WRITE | NOTE_EXTEND | NOTE_DELETE | NOTE_RENAME; + KJ_SYSCALL(kevent(kqueueFd, &change, 1, nullptr, 0, nullptr)); + filesWatched.add(kj::mv(fd)); + } +}; + +#else + +// Dummy backend for platforms without an implementation (Windows, ...), mirroring workerd's: +// isSupported() returns false, which workerd surfaces as a clean CLI error for --watch +// ("File watching is not yet implemented on your OS") rather than a crash. A real Windows +// backend (e.g. ReadDirectoryChangesW, perhaps via the notify crate) is a potential follow-up. +struct FileWatcher::Impl { + bool isSupported() { + return false; + } + + void watch(kj::PathPtr path, kj::Maybe file) {} + + kj::Promise onChange() { + return kj::NEVER_DONE; + } +}; + +#endif + +FileWatcher::FileWatcher(): impl(kj::heap()) {} +FileWatcher::~FileWatcher() noexcept(false) = default; + +bool FileWatcher::isSupported() { + return impl->isSupported(); +} + +void FileWatcher::watch(kj::PathPtr path, kj::Maybe file) { + impl->watch(path, file); +} + +kj::Promise FileWatcher::onChange() { + return impl->onChange(); +} + +} // namespace kj_rs_io diff --git a/src/rust/cxx/kj-rs-io/file-watcher.h b/src/rust/cxx/kj-rs-io/file-watcher.h new file mode 100644 index 00000000000..6b2a85b2fe6 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/file-watcher.h @@ -0,0 +1,54 @@ +#pragma once +// kj_rs_io::FileWatcher: the tokio-loop replacement for workerd's `--watch` file watcher. +// +// Watches a set of individual files and resolves onChange() when any of them changes. The +// platform backends mirror workerd's kj FileWatcher exactly — inotify on the parent directory +// on Linux, kqueue EVFILT_VNODE per open file on macOS/BSD — but the notification fd's +// readiness is awaited through tokio's AsyncFd (kj_rs_io::wait_fd_readable) instead of +// kj::UnixEventPort::FdObserver, so it works on the tokio-backed event loop where no +// UnixEventPort exists. +// +// Behavior notes (all matching the kj version): +// - Multiple rapid changes coalesce: onChange() resolves once for whatever is queued; calling +// it again drains the queue before waiting, so changes are never lost between calls. +// - Linux: the watched file itself need not exist (only its parent directory must), and a +// file deleted and re-created is picked up again. macOS/BSD: watch() opens the file (or +// dups the provided already-open fd), so watching a nonexistent file throws; a +// replaced-by-rename file still fires on the old inode. +// - All internal fds are CLOEXEC: --watch reloads via execve(), which must not leak them. +// - Unsupported platforms (Windows, ...): isSupported() returns false, watch() is a no-op and +// onChange() never resolves, mirroring workerd's dummy watcher. +// +// onChange() must be awaited on the thread owning the kj_rs_tokio::TokioEventPort, and at most +// one onChange() promise may be outstanding at a time (workerd awaits it sequentially). The +// promise must be dropped before the FileWatcher is destroyed. + +#include +#include + +namespace kj_rs_io { + +class FileWatcher { + public: + FileWatcher(); + ~FileWatcher() noexcept(false); + KJ_DISALLOW_COPY_AND_MOVE(FileWatcher); + + // False on platforms with no watcher implementation (callers should report an error). + bool isSupported(); + + // Adds `path` to the watched set. `file`, if provided, is an already-open handle for the + // same path (the kqueue backend watches it directly via a dup'd fd; the inotify backend + // ignores it and watches the parent directory by name). + void watch(kj::PathPtr path, kj::Maybe file); + + // Resolves the next time any watched file changes (immediately, if a change is already + // queued). Eagerly evaluated, per kj-rs-io convention for I/O promises. + kj::Promise onChange(); + + private: + struct Impl; + kj::Own impl; +}; + +} // namespace kj_rs_io diff --git a/src/rust/cxx/kj-rs-io/lib.rs b/src/rust/cxx/kj-rs-io/lib.rs new file mode 100644 index 00000000000..f9a071b6c5c --- /dev/null +++ b/src/rust/cxx/kj-rs-io/lib.rs @@ -0,0 +1,75 @@ +//! Rust half of kj-rs-io: tokio-backed implementations of KJ's async I/O interfaces. +//! +//! The C++ side (`async-io.h`) implements `kj::AsyncIoStream`, `kj::ConnectionReceiver`, +//! `kj::NetworkAddress`, `kj::Network`, `kj::AsyncIoProvider` and `kj::LowLevelAsyncIoProvider` +//! as thin wrappers over the opaque Rust types in this crate. All async operations are plain +//! `async fn`s bridged to `kj::Promise` by workerd-cxx; dropping the promise drops the Rust +//! future, which releases any tokio readiness interest (cancellation is implicit). +//! +//! Every future returned from this crate must be polled on the thread that owns the +//! `kj_rs_tokio::TokioEventPort` runtime: tokio I/O objects register with that runtime's I/O +//! driver, which is only driven while the KJ event loop sleeps inside the port's +//! `wait()`/`poll()`. +//! +//! A Rust-originated stream wrapped as `kj::AsyncIoStream` can be recovered as its native +//! tokio object so Rust servers can drive the connection without crossing the +//! FFI per read: [`unwrap_kj_stream`] from Rust, `kj_rs_io::unwrapTokioStream()` from C++. +//! Foreign streams fail to unwrap with a `kj::Exception`. The native-serve entry points +//! ([`serve_kj_stream`], [`take_kj_socket`]) build on the [`serve`] module's `ServeIo` / pump +//! machinery. + +// Safety & panic enforcement walls. Test code exempted. +// +// `unsafe` is quarantined into a single named FFI island: the crate root denies `unsafe_code`, so +// the serve / net / stream / error / runtime / readiness / signal business logic is +// *compiler-proven* free of hand-written unsafe. The one island that opts back in via +// `#![allow(unsafe_code)]` is `ffi.rs`: it holds the `#[cxx::bridge] mod bridge` (re-exported as +// `crate::ffi::*`) plus all the fd / sockaddr laundering and the one remaining `pub unsafe fn` +// FFI entry point (`unwrap_kj_stream` — borrow-based, C++ keeps the wrapper). The owning entry +// points (`take_kj_socket`, `serve_kj_stream`) are safe fns in `serve.rs`: ownership arrives as +// a `KjOwn`, the pump drives it through compiler-checked shared borrows (shared-receiver shims, +// see unwrap.h), and no raw pointer crosses the public surface. +#![deny(unsafe_op_in_unsafe_fn)] +#![deny(unsafe_code)] +#![deny(clippy::undocumented_unsafe_blocks)] +#![deny( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::unreachable, + clippy::todo, + clippy::unimplemented +)] +#![cfg_attr( + test, + allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::unreachable, + clippy::todo, + clippy::unimplemented + ) +)] + +pub use stream::TokioStream; + +mod error; +mod ffi; +mod net; +mod readiness; +mod runtime; +pub mod serve; +mod signal; +mod stream; + +/// Opaque binding of `kj::AsyncIoStream` (see [`unwrap_kj_stream`]). +pub use ffi::KjAsyncIoStream; +pub use ffi::unwrap_kj_stream; +pub use serve::ServeIo; +pub use serve::ServePath; +pub use serve::ServedKjStream; +pub use serve::StreamPump; +pub use serve::TakeSocketError; +pub use serve::serve_kj_stream; +pub use serve::take_kj_socket; diff --git a/src/rust/cxx/kj-rs-io/net.rs b/src/rust/cxx/kj-rs-io/net.rs new file mode 100644 index 00000000000..c2a83c18f91 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/net.rs @@ -0,0 +1,631 @@ +//! Tokio-backed `kj::Network` / `kj::NetworkAddress` / `kj::ConnectionReceiver` backends. +//! +//! Address-string grammar follows KJ's `SocketAddress::parse` (kj/async-io-unix.c++) for the +//! subset workerd feeds it: +//! +//! - IPv4: `"1.2.3.4"`, `"1.2.3.4:80"` +//! - IPv6: `"1234:5678::abcd"`, `"[1234:5678::abcd]:80"` +//! - Wildcard (dual-stack): `"*"`, `"*:80"` +//! - Hostnames (DNS via blocking `getaddrinfo` on tokio's blocking pool, its completion delivered +//! same-thread through a tokio runtime task — see [`resolve_host`]): `"example.com"`, +//! `"example.com:80"` +//! - Unix domain: `"unix:/path/to/socket"` (Unix only) +//! +//! Known deviations from KJ, all erroring loudly rather than misbehaving: named services +//! (`"host:http"`), `unix-abstract:` addresses, and IPv6 scope IDs (`"fe80::1%eth0"`) are not +//! supported. + +use std::net::IpAddr; +use std::net::SocketAddr; +use std::net::ToSocketAddrs; + +use tokio::net::TcpListener; +use tokio::net::TcpStream; +#[cfg(unix)] +use tokio::net::UnixListener; +#[cfg(unix)] +use tokio::net::UnixStream; + +use crate::error::KjIoError; +use crate::error::Result; +use crate::error::op; +use crate::runtime::runtime_handle; +use crate::runtime::with_runtime; +use crate::stream::TokioStream; + +const LISTEN_BACKLOG: i32 = 1024; + +/// A parsed network address: one or more socket addresses to try in order. +pub struct TokioAddress { + spec: Spec, +} + +#[derive(Clone)] +enum Spec { + Ip { + /// Resolved addresses, tried in order by `connect()`; `listen()` binds the first one + /// (mirroring KJ, which also only listens on the first result). + addrs: Vec, + /// `"*"`: listen on `[::]` with `IPV6_V6ONLY` disabled (dual-stack), reject `connect()`. + wildcard: bool, + }, + #[cfg(unix)] + Unix { path: std::path::PathBuf }, +} + +impl TokioAddress { + async fn parse(text: &str, port_hint: u16) -> Result { + if let Some(path) = text.strip_prefix("unix:") { + #[cfg(unix)] + { + return Ok(Self { + spec: Spec::Unix { path: path.into() }, + }); + } + #[cfg(not(unix))] + { + let _ = path; + return Err(KjIoError::other( + "parseAddress", + "Unix domain sockets are not supported on this platform", + )); + } + } + if text.starts_with("unix-abstract:") { + return Err(KjIoError::other( + "parseAddress", + "abstract Unix domain sockets are not implemented by kj-rs-io", + )); + } + + // Split into address and port parts, exactly like KJ's SocketAddress::parse. + let (addr_part, port_part) = if let Some(rest) = text.strip_prefix('[') { + // Bracketed IPv6, optionally "[..]:port". + let close = rest.rfind(']').ok_or_else(|| { + KjIoError::other("parseAddress", format!("Unclosed '[' in address: {text}")) + })?; + let addr = &rest[..close]; + let tail = &rest[close + 1..]; + if tail.is_empty() { + (addr, None) + } else if let Some(port) = tail.strip_prefix(':') { + (addr, Some(port)) + } else { + return Err(KjIoError::other( + "parseAddress", + format!("Expected port suffix after ']': {text}"), + )); + } + } else if let Some(colon) = text.find(':') { + if text[colon + 1..].contains(':') { + // Two or more colons, no brackets: a bare IPv6 address with no port. + (text, None) + } else { + // Exactly one colon: ip4/hostname with port. + (&text[..colon], Some(&text[colon + 1..])) + } + } else { + (text, None) + }; + + let port = match port_part { + Some(port_text) => port_text.parse::().map_err(|_| { + // KJ falls back to getaddrinfo service-name resolution here; tokio's resolver + // only accepts numeric ports. + KjIoError::other( + "parseAddress", + format!("invalid port (named services are not supported): {port_text}"), + ) + })?, + None => port_hint, + }; + + if addr_part == "*" { + return Ok(Self { + spec: Spec::Ip { + addrs: vec![SocketAddr::new( + IpAddr::V6(std::net::Ipv6Addr::UNSPECIFIED), + port, + )], + wildcard: true, + }, + }); + } + + if let Ok(ip) = addr_part.parse::() { + return Ok(Self { + spec: Spec::Ip { + addrs: vec![SocketAddr::new(ip, port)], + wildcard: false, + }, + }); + } + + // Not a literal: resolve the hostname via getaddrinfo. To honor kj-rs's single-thread + // axiom, the blocking getaddrinfo runs on tokio's blocking pool but its completion is + // absorbed by tokio's own scheduler (via a runtime task) and forwarded on the loop thread, + // so this await resumes same-thread (never cross-thread). See `resolve_host`. + let addrs: Vec = resolve_host(addr_part, port).await?; + if addrs.is_empty() { + return Err(KjIoError::other( + "getaddrinfo()", + format!("no addresses found for host: {addr_part}"), + )); + } + Ok(Self { + spec: Spec::Ip { + addrs, + wildcard: false, + }, + }) + } + + async fn connect_index(&self, index: usize) -> Result> { + match &self.spec { + Spec::Ip { addrs, wildcard } => { + if *wildcard { + return Err(KjIoError::other( + "connect()", + "cannot connect() to a wildcard address", + )); + } + let addr = addrs + .get(index) + .ok_or_else(|| KjIoError::other("connect()", "address index out of range"))?; + let stream = TcpStream::connect(addr).await.map_err(op("connect()"))?; + Ok(Box::new(TokioStream::from_tcp(stream))) + } + #[cfg(unix)] + Spec::Unix { path } => { + if index != 0 { + return Err(KjIoError::other("connect()", "address index out of range")); + } + let stream = UnixStream::connect(path).await.map_err(op("connect()"))?; + Ok(Box::new(TokioStream::from_unix(stream))) + } + } + } + + fn count(&self) -> usize { + match &self.spec { + Spec::Ip { addrs, .. } => addrs.len(), + #[cfg(unix)] + Spec::Unix { .. } => 1, + } + } + + fn raw_sockaddr(&self, index: usize) -> Result> { + let sockaddr: socket2::SockAddr = match &self.spec { + Spec::Ip { addrs, .. } => (*addrs + .get(index) + .ok_or_else(|| KjIoError::other("sockaddr", "address index out of range"))?) + .into(), + #[cfg(unix)] + Spec::Unix { path } => { + if index != 0 { + return Err(KjIoError::other("sockaddr", "address index out of range")); + } + socket2::SockAddr::unix(path).map_err(op("sockaddr"))? + } + }; + Ok(crate::ffi::sockaddr_to_bytes(&sockaddr)) + } + + fn listen(&self) -> Result> { + let handle = runtime_handle()?; + match &self.spec { + Spec::Ip { addrs, wildcard } => { + let addr = *addrs + .first() + .ok_or_else(|| KjIoError::other("listen()", "no addresses to bind"))?; + let domain = socket2::Domain::for_address(addr); + let socket = socket2::Socket::new(domain, socket2::Type::STREAM, None) + .map_err(op("socket()"))?; + // KJ parity: SO_REUSEADDR on listeners; wildcard sockets accept both address + // families (IPV6_V6ONLY off). + socket.set_reuse_address(true).map_err(op("setsockopt()"))?; + if *wildcard { + socket.set_only_v6(false).map_err(op("setsockopt()"))?; + } + socket.bind(&addr.into()).map_err(op("bind()"))?; + socket.listen(LISTEN_BACKLOG).map_err(op("listen()"))?; + socket.set_nonblocking(true).map_err(op("fcntl()"))?; + // Registering with the I/O driver requires the runtime context. + let _guard = handle.enter(); + let listener = TcpListener::from_std(socket.into()).map_err(op("wrap listener"))?; + Ok(Box::new(TokioListener { + inner: ListenerInner::Tcp(listener), + })) + } + #[cfg(unix)] + Spec::Unix { path } => { + // Like KJ, no unlink(): binding an existing path fails. + let listener = + std::os::unix::net::UnixListener::bind(path).map_err(op("bind()"))?; + listener.set_nonblocking(true).map_err(op("fcntl()"))?; + let _guard = handle.enter(); + let listener = UnixListener::from_std(listener).map_err(op("wrap listener"))?; + Ok(Box::new(TokioListener { + inner: ListenerInner::Unix(listener), + })) + } + } + } + + fn to_display_string(&self) -> String { + match &self.spec { + Spec::Ip { addrs, wildcard } => { + if *wildcard { + format!("*:{}", addrs[0].port()) + } else { + let parts: Vec = addrs.iter().map(ToString::to_string).collect(); + parts.join(",") + } + } + #[cfg(unix)] + Spec::Unix { path } => format!("unix:{}", path.display()), + } + } +} + +/// A listening socket (`kj::ConnectionReceiver` backend). +pub struct TokioListener { + inner: ListenerInner, +} + +enum ListenerInner { + Tcp(TcpListener), + #[cfg(unix)] + Unix(UnixListener), +} + +impl TokioListener { + async fn accept(&self) -> Result> { + match &self.inner { + ListenerInner::Tcp(listener) => { + let (stream, _peer) = listener.accept().await.map_err(op("accept()"))?; + let _ = stream.set_nodelay(true); + Ok(Box::new(TokioStream::from_tcp(stream))) + } + #[cfg(unix)] + ListenerInner::Unix(listener) => { + let (stream, _peer) = listener.accept().await.map_err(op("accept()"))?; + Ok(Box::new(TokioStream::from_unix(stream))) + } + } + } + + fn port(&self) -> Result { + match &self.inner { + ListenerInner::Tcp(listener) => { + Ok(listener.local_addr().map_err(op("getsockname()"))?.port()) + } + // KJ returns 0 for non-IP listeners. + #[cfg(unix)] + ListenerInner::Unix(_) => Ok(0), + } + } + + /// Borrows the live listener socket's fd (tokio listeners implement `AsFd`), for the + /// sockopt/sockname passthrough behind `kj::ConnectionReceiver`. + #[cfg(unix)] + pub(crate) fn as_borrowed_fd(&self) -> std::os::fd::BorrowedFd<'_> { + use std::os::fd::AsFd; + match &self.inner { + ListenerInner::Tcp(listener) => listener.as_fd(), + ListenerInner::Unix(listener) => listener.as_fd(), + } + } + + /// Borrows the live listener socket's `SOCKET` (tokio's `TcpListener` implements + /// `AsSocket`): the Windows counterpart of [`TokioListener::as_borrowed_fd`]. On Windows + /// only the Tcp variant of `ListenerInner` exists. + // Validated by Windows CI; mirrors the unix arm. + #[cfg(windows)] + pub(crate) fn as_borrowed_socket(&self) -> std::os::windows::io::BorrowedSocket<'_> { + use std::os::windows::io::AsSocket; + match &self.inner { + ListenerInner::Tcp(listener) => listener.as_socket(), + } + } + + /// Raw `struct sockaddr` bytes of the listener's bound address (the `getsockname()` + /// passthrough behind `kj::ConnectionReceiver::getsockname`). + #[cfg(any(unix, windows))] + fn local_addr_bytes(&self) -> Result> { + #[cfg(unix)] + let sock = self.as_borrowed_fd(); + // Validated by Windows CI; mirrors the unix arm. + #[cfg(windows)] + let sock = self.as_borrowed_socket(); + let addr = socket2::SockRef::from(&sock) + .local_addr() + .map_err(op("getsockname()"))?; + Ok(crate::ffi::sockaddr_to_bytes(&addr)) + } + + #[cfg(not(any(unix, windows)))] + fn local_addr_bytes(&self) -> Result> { + Err(KjIoError::other( + "getsockname", + "not implemented by kj-rs-io on this platform", + )) + } +} + +// ====================================================================================== +// Bridge entry points (see lib.rs). + +/// Resolves a hostname via blocking `getaddrinfo` on tokio's blocking pool — purely in Rust and +/// same-thread. A tokio *runtime* task owns the blocking `JoinHandle`, so the blocking-pool +/// completion wakes tokio's own (`Send + Sync`) scheduler waker — which unparks this loop — +/// rather than the bridged future's kj waker. The task then runs on the loop thread and hands the +/// result back over a oneshot, waking the awaiting future same-thread. No cross-thread rust waker +/// is involved (that is the whole reason we do not use `tokio::net::lookup_host`, whose +/// `JoinHandle` wake lands on the caller's waker cross-thread). Same blocking `getaddrinfo` call +/// and NSS/`/etc/hosts` parity as `lookup_host`. +async fn resolve_host(host: &str, port: u16) -> Result> { + let host = host.to_owned(); + let (tx, rx) = tokio::sync::oneshot::channel::>>(); + + // The runtime task's waker is tokio's own scheduler waker (Send + Sync), so the cross-thread + // completion from the blocking pool terminates inside tokio's scheduler (unparking this loop), + // never at a rust cross-thread waker. The task forwards the result on the loop thread, waking + // the awaiting future same-thread. + let task = tokio::spawn(async move { + let resolved = match tokio::task::spawn_blocking(move || { + (host.as_str(), port) + .to_socket_addrs() + .map(std::iter::Iterator::collect::>) + }) + .await + { + Ok(result) => result, + Err(_) => Err(std::io::Error::other("getaddrinfo task failed")), + }; + let _ = tx.send(resolved); + }); + // If this future is dropped (KJ promise cancelled), abort the forwarding task rather than + // leaving it to run to completion for nobody. The blocking getaddrinfo call itself cannot be + // interrupted once started (an OS limitation shared with KJ's own resolver and tokio's + // `lookup_host`), so its blocking-pool slot is reclaimed only when the syscall returns. + let _abort_guard = crate::runtime::AbortOnDrop(task); + + match rx.await { + Ok(result) => result.map_err(op("getaddrinfo()")), + Err(_) => Err(KjIoError::other( + "getaddrinfo()", + "DNS resolver task dropped", + )), + } +} + +pub async fn network_parse_address(addr: String, port_hint: u16) -> Result> { + with_runtime(async move { Ok(Box::new(TokioAddress::parse(&addr, port_hint).await?)) }).await +} + +pub fn network_get_sockaddr(sockaddr: &[u8]) -> Result> { + #[cfg(any(unix, windows))] + { + let addr = sockaddr_from_bytes(sockaddr)?; + if let Some(socket_addr) = addr.as_socket() { + return Ok(Box::new(TokioAddress { + spec: Spec::Ip { + addrs: vec![socket_addr], + wildcard: false, + }, + })); + } + #[cfg(unix)] + if let Some(path) = addr.as_pathname() { + return Ok(Box::new(TokioAddress { + spec: Spec::Unix { path: path.into() }, + })); + } + Err(KjIoError::other( + "getSockaddr", + "unsupported sockaddr family", + )) + } + #[cfg(not(any(unix, windows)))] + { + let _ = sockaddr; + Err(KjIoError::other( + "getSockaddr", + "not implemented on this platform", + )) + } +} + +/// Connects to exactly the `index`th resolved address (no fallback). The C++ side drives the +/// try-each-address loop itself so it can apply `restrictPeers()` filtering per address before +/// initiating each connection attempt (KJ parity: a blocked address contributes a +/// "`connect()` blocked by `restrictPeers()`" failure; only the last address's error propagates). +pub async fn address_connect_index(addr: &TokioAddress, index: usize) -> Result> { + with_runtime(addr.connect_index(index)).await +} + +/// Number of resolved socket addresses behind this address (>= 1). +pub fn address_count(addr: &TokioAddress) -> usize { + addr.count() +} + +/// Raw `struct sockaddr` bytes of the `index`th resolved address, for the C++ side's +/// `kj::_::NetworkFilter` (restrictPeers) checks. +pub fn address_raw_sockaddr(addr: &TokioAddress, index: usize) -> Result> { + addr.raw_sockaddr(index) +} + +pub fn address_listen(addr: &TokioAddress) -> Result> { + addr.listen() +} + +#[expect(clippy::unnecessary_box_returns)] // Opaque cxx types must cross the bridge boxed. +pub fn address_clone(addr: &TokioAddress) -> Box { + Box::new(TokioAddress { + spec: addr.spec.clone(), + }) +} + +pub fn address_to_string(addr: &TokioAddress) -> String { + addr.to_display_string() +} + +pub async fn listener_accept(listener: &TokioListener) -> Result> { + with_runtime(listener.accept()).await +} + +pub fn listener_port(listener: &TokioListener) -> Result { + listener.port() +} + +pub fn listener_local_addr(listener: &TokioListener) -> Result> { + listener.local_addr_bytes() +} + +// ====================================================================================== +// Socket-handle wrapping. All handles arrive owned and non-blocking as an `i64` "raw socket +// handle" — a Unix fd or a win32 SOCKET (the C++ side normalizes KJ's TAKE_OWNERSHIP / +// ALREADY_CLOEXEC / ALREADY_NONBLOCK flags, dup'ing when not taking ownership). The platform +// split lives entirely in `ffi::own_socket_from_raw` (the one conversion point); everything +// here operates on the uniform `socket2::Socket` / std / tokio types. + +#[cfg(any(unix, windows))] +fn socket_from_raw(handle: i64) -> socket2::Socket { + crate::ffi::own_socket_from_raw(handle) +} + +#[cfg(any(unix, windows))] +fn sockaddr_from_bytes(bytes: &[u8]) -> Result { + crate::ffi::sockaddr_from_bytes(bytes) +} + +/// The handle tier of [`crate::take_kj_socket`] (unix only; see that function's docs): wraps +/// an *owned*, connected stream-socket fd (TCP or Unix domain, detected automatically) as a +/// [`crate::serve::ServeIo`]. Unlike [`wrap_socket_fd`] the fd is a fresh dup of a kj stream's +/// socket, so non-blocking mode is forced rather than assumed (the original may have come from +/// anywhere). +#[cfg(unix)] +pub fn serve_io_from_owned_fd(fd: std::os::fd::OwnedFd) -> Result { + let socket = socket2::Socket::from(fd); + socket.set_nonblocking(true).map_err(op("fcntl()"))?; + let local = socket.local_addr().map_err(op("getsockname()"))?; + let _guard = runtime_handle()?.enter(); + match local.domain() { + socket2::Domain::IPV4 | socket2::Domain::IPV6 => { + let stream = TcpStream::from_std(socket.into()).map_err(op("takeKjSocket"))?; + Ok(crate::serve::ServeIo::Tcp(stream)) + } + socket2::Domain::UNIX => { + let stream = UnixStream::from_std(socket.into()).map_err(op("takeKjSocket"))?; + Ok(crate::serve::ServeIo::Unix(stream)) + } + _ => Err(KjIoError::other( + "takeKjSocket", + "unsupported socket family", + )), + } +} + +pub fn wrap_socket_fd(handle: i64) -> Result> { + #[cfg(any(unix, windows))] + { + let socket = socket_from_raw(handle); + let local = socket.local_addr().map_err(op("getsockname()"))?; + let _guard = runtime_handle()?.enter(); + match local.domain() { + socket2::Domain::IPV4 | socket2::Domain::IPV6 => { + let stream = TcpStream::from_std(socket.into()).map_err(op("wrapSocketFd"))?; + Ok(Box::new(TokioStream::from_tcp(stream))) + } + #[cfg(unix)] + socket2::Domain::UNIX => { + let stream = UnixStream::from_std(socket.into()).map_err(op("wrapSocketFd"))?; + Ok(Box::new(TokioStream::from_unix(stream))) + } + _ => Err(KjIoError::other( + "wrapSocketFd", + "unsupported socket family", + )), + } + } + #[cfg(not(any(unix, windows)))] + { + let _ = handle; + Err(KjIoError::other( + "wrapSocketFd", + "not implemented on this platform", + )) + } +} + +pub fn wrap_listen_fd(handle: i64) -> Result> { + #[cfg(any(unix, windows))] + { + let socket = socket_from_raw(handle); + let local = socket.local_addr().map_err(op("getsockname()"))?; + let _guard = runtime_handle()?.enter(); + match local.domain() { + socket2::Domain::IPV4 | socket2::Domain::IPV6 => { + let listener = + TcpListener::from_std(socket.into()).map_err(op("wrapListenSocketFd"))?; + Ok(Box::new(TokioListener { + inner: ListenerInner::Tcp(listener), + })) + } + #[cfg(unix)] + socket2::Domain::UNIX => { + let listener = + UnixListener::from_std(socket.into()).map_err(op("wrapListenSocketFd"))?; + Ok(Box::new(TokioListener { + inner: ListenerInner::Unix(listener), + })) + } + _ => Err(KjIoError::other( + "wrapListenSocketFd", + "unsupported socket family", + )), + } + } + #[cfg(not(any(unix, windows)))] + { + let _ = handle; + Err(KjIoError::other( + "wrapListenSocketFd", + "not implemented on this platform", + )) + } +} + +pub async fn wrap_connecting_socket_fd(handle: i64, sockaddr: Vec) -> Result> { + #[cfg(any(unix, windows))] + { + with_runtime(async move { + let addr = sockaddr_from_bytes(&sockaddr)?; + let socket_addr = addr.as_socket().ok_or_else(|| { + KjIoError::other( + "wrapConnectingSocketFd", + "only AF_INET/AF_INET6 sockaddrs are supported", + ) + })?; + let socket = socket_from_raw(handle); + // TcpSocket::connect handles the nonblocking connect dance (EINPROGRESS, wait for + // writability, check SO_ERROR) and registers with the I/O driver. + let tcp_socket = tokio::net::TcpSocket::from_std_stream(socket.into()); + let stream = tcp_socket + .connect(socket_addr) + .await + .map_err(op("connect()"))?; + Ok(Box::new(TokioStream::from_tcp(stream))) + }) + .await + } + #[cfg(not(any(unix, windows)))] + { + let _ = (handle, sockaddr); + Err(KjIoError::other( + "wrapConnectingSocketFd", + "not implemented on this platform", + )) + } +} diff --git a/src/rust/cxx/kj-rs-io/peer-filter.c++ b/src/rust/cxx/kj-rs-io/peer-filter.c++ new file mode 100644 index 00000000000..61a3ccce563 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/peer-filter.c++ @@ -0,0 +1,201 @@ +// Port of KJ's kj::_::NetworkFilter (kj/async-io.c++, MIT-licensed, Sandstorm Development +// Group and contributors) — see peer-filter.h for why this is a port rather than a reuse. +// Behavior must be kept in lockstep with upstream KJ. + +#include "kj-rs-io/peer-filter.h" + +#include + +#if _WIN32 +#include +#else +#include +#include +#include +#endif + +namespace kj_rs_io { +namespace { + +using kj::CidrRange; + +kj::ArrayPtr localCidrs() { + static const CidrRange result[] = { + // localhost + "127.0.0.0/8"_kj, + "::1/128"_kj, + + // Trying to *connect* to 0.0.0.0 on many systems is equivalent to connecting to + // localhost. (wat) + "0.0.0.0/32"_kj, + "::/128"_kj, + }; + return kj::arrayPtr(result, kj::size(result)); +} + +kj::ArrayPtr privateCidrs() { + static const CidrRange result[] = { + "10.0.0.0/8"_kj, // RFC1918 reserved for internal network + "100.64.0.0/10"_kj, // RFC6598 "shared address space" for carrier-grade NAT + "169.254.0.0/16"_kj, // RFC3927 "link local" (auto-configured LAN in absence of DHCP) + "172.16.0.0/12"_kj, // RFC1918 reserved for internal network + "192.168.0.0/16"_kj, // RFC1918 reserved for internal network + + "fc00::/7"_kj, // RFC4193 unique private network + "fe80::/10"_kj, // RFC4291 "link local" (auto-configured LAN in absence of DHCP) + }; + return kj::arrayPtr(result, kj::size(result)); +} + +kj::ArrayPtr reservedCidrs() { + // Address ranges reserved by RFCs for specific alternative protocols. These are not + // considered part of "public", "private", "network", nor "local". But, we will allow apps to + // explicitly allowlist CIDRs in this range if they really want, because some people actually + // use these ranges as if they were private ranges. + static const CidrRange result[] = { + "192.0.0.0/24"_kj, // RFC6890 reserved for special protocols + "224.0.0.0/4"_kj, // RFC1112 multicast + "240.0.0.0/4"_kj, // RFC1112 multicast / reserved for future use + "255.255.255.255/32"_kj, // RFC0919 broadcast address + + "2001::/23"_kj, // RFC2928 reserved for special protocols + "ff00::/8"_kj, // RFC4291 multicast + }; + return kj::arrayPtr(result, kj::size(result)); +} + +bool matchesAny(kj::ArrayPtr cidrs, const struct sockaddr *addr) { + for (auto &cidr: cidrs) { + if (cidr.matches(addr)) return true; + } + return false; +} + +#if !_WIN32 +// sockaddr_un::sun_path is not required to have a NUL terminator, so it must be read carefully. +kj::ArrayPtr safeUnixPath(const struct sockaddr_un *addr, kj::uint addrlen) { + KJ_REQUIRE(addr->sun_family == AF_UNIX, "not a unix address"); + KJ_REQUIRE(addrlen >= offsetof(sockaddr_un, sun_path), "invalid unix address"); + + size_t maxPathlen = addrlen - offsetof(sockaddr_un, sun_path); + + size_t pathlen; + if (maxPathlen > 0 && addr->sun_path[0] == '\0') { + // Linux "abstract" unix address + pathlen = strnlen(addr->sun_path + 1, maxPathlen - 1) + 1; + } else { + pathlen = strnlen(addr->sun_path, maxPathlen); + } + return kj::arrayPtr(addr->sun_path, pathlen); +} +#endif // !_WIN32 + +} // namespace + +PeerFilter::PeerFilter(): allowUnix(true), allowAbstractUnix(true) { + allowCidrs.add(CidrRange::inet4({0, 0, 0, 0}, 0)); + allowCidrs.add(CidrRange::inet6({}, {}, 0)); +} + +PeerFilter::PeerFilter(kj::ArrayPtr allow, + kj::ArrayPtr deny, + PeerFilter &next) + : allowUnix(false), + allowAbstractUnix(false), + next(next) { + for (auto rule: allow) { + if (rule == "local") { + allowCidrs.addAll(localCidrs()); + } else if (rule == "network") { + // Can't be represented as a simple union of CIDRs, so we handle in shouldAllow(). + allowNetwork = true; + } else if (rule == "private") { + allowCidrs.addAll(privateCidrs()); + allowCidrs.addAll(localCidrs()); + } else if (rule == "public") { + // Can't be represented as a simple union of CIDRs, so we handle in shouldAllow(). + allowPublic = true; + } else if (rule == "unix") { + allowUnix = true; + } else if (rule == "unix-abstract") { + allowAbstractUnix = true; + } else { + allowCidrs.add(CidrRange(rule)); + } + } + + for (auto rule: deny) { + if (rule == "local") { + denyCidrs.addAll(localCidrs()); + } else if (rule == "network") { + KJ_FAIL_REQUIRE("don't deny 'network', allow 'local' instead"); + } else if (rule == "private") { + denyCidrs.addAll(privateCidrs()); + } else if (rule == "public") { + // Tricky: What if we allow 'network' and deny 'public'? + KJ_FAIL_REQUIRE("don't deny 'public', allow 'private' instead"); + } else if (rule == "unix") { + allowUnix = false; + } else if (rule == "unix-abstract") { + allowAbstractUnix = false; + } else { + denyCidrs.add(CidrRange(rule)); + } + } +} + +bool PeerFilter::shouldAllow(const struct sockaddr *addr, kj::uint addrlen) { + KJ_REQUIRE(addrlen >= sizeof(addr->sa_family)); + +#if !_WIN32 + if (addr->sa_family == AF_UNIX) { + auto path = safeUnixPath(reinterpret_cast(addr), addrlen); + if (path.size() > 0 && path[0] == '\0') { + return allowAbstractUnix; + } else { + return allowUnix; + } + } +#endif + + bool allowed = false; + kj::uint allowSpecificity = 0; + + if (allowPublic) { + if ((addr->sa_family == AF_INET || addr->sa_family == AF_INET6) && + !matchesAny(privateCidrs(), addr) && !matchesAny(localCidrs(), addr) && + !matchesAny(reservedCidrs(), addr)) { + allowed = true; + // Don't adjust allowSpecificity as this match has an effective specificity of zero. + } + } + + if (allowNetwork) { + if ((addr->sa_family == AF_INET || addr->sa_family == AF_INET6) && + !matchesAny(localCidrs(), addr) && !matchesAny(reservedCidrs(), addr)) { + allowed = true; + // Don't adjust allowSpecificity as this match has an effective specificity of zero. + } + } + + for (auto &cidr: allowCidrs) { + if (cidr.matches(addr)) { + allowSpecificity = kj::max(allowSpecificity, cidr.getSpecificity()); + allowed = true; + } + } + if (!allowed) return false; + for (auto &cidr: denyCidrs) { + if (cidr.matches(addr)) { + if (cidr.getSpecificity() >= allowSpecificity) return false; + } + } + + KJ_IF_SOME(n, next) { + return n.shouldAllow(addr, addrlen); + } else { + return true; + } +} + +} // namespace kj_rs_io diff --git a/src/rust/cxx/kj-rs-io/peer-filter.h b/src/rust/cxx/kj-rs-io/peer-filter.h new file mode 100644 index 00000000000..9745acce2b3 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/peer-filter.h @@ -0,0 +1,54 @@ +#pragma once +// PeerFilter: a faithful port of KJ's kj::_::NetworkFilter (kj/async-io.c++), backing +// kj-rs-io's Network::restrictPeers() support. +// +// Ported rather than reused because kj::_::NetworkFilter lives in KJ's internal header +// (kj/async-io-internal.h), whose quoted includes ("vector.h") only resolve inside the KJ +// source tree — it is not includable through Bazel's virtual include dirs. The allow/deny +// grammar ("public"/"private"/"local"/"network"/"unix"/"unix-abstract"/CIDRs), the RFC CIDR +// tables, the specificity tie-breaking between allow and deny rules, and the filter-chaining +// semantics are kept identical so restrictPeers() behaves exactly like kj::setupAsyncIo()'s +// networks. kj::CidrRange itself IS reused (kj/cidr.h is a clean public header). + +#include +#include +#include + +namespace kj_rs_io { + +class PeerFilter final: public kj::LowLevelAsyncIoProvider::NetworkFilter { + public: + // Allow-everything filter (matches KJ's root networks). + PeerFilter(); + + // Restriction layered on `next` (which must outlive this filter). Grammar identical to + // kj::Network::restrictPeers(). + PeerFilter(kj::ArrayPtr allow, + kj::ArrayPtr deny, + PeerFilter &next); + + // Read-only despite the non-const signature: this override matches + // kj::LowLevelAsyncIoProvider::NetworkFilter::shouldAllow (declared non-const upstream), but the + // implementation only *reads* the CIDR tables / flags and recurses into `next` — it mutates no + // member and has no interior mutability, so concurrent callers sharing a filter are safe. Keep + // it read-only. + bool shouldAllow(const struct sockaddr *addr, kj::uint addrlen) override; + + // Immobile like KJ's own kj::_::NetworkFilter: a restricted filter's `next` points at a parent + // filter, and derived (restrictPeers) filters point back at this one, so a move would dangle the + // chain. Every instance is either an owning member of a heap-allocated network/provider or + // kj::heap(), so nothing moves one; this guards the latent hazard. + KJ_DISALLOW_COPY_AND_MOVE(PeerFilter); + + private: + kj::Vector allowCidrs; + kj::Vector denyCidrs; + bool allowUnix; + bool allowAbstractUnix; + bool allowPublic = false; + bool allowNetwork = false; + + kj::Maybe next; +}; + +} // namespace kj_rs_io diff --git a/src/rust/cxx/kj-rs-io/readiness.rs b/src/rust/cxx/kj-rs-io/readiness.rs new file mode 100644 index 00000000000..63b6202e67a --- /dev/null +++ b/src/rust/cxx/kj-rs-io/readiness.rs @@ -0,0 +1,49 @@ +//! tokio-backed fd readiness watching. +//! +//! This backs `kj_rs_io::FileWatcher` (file-watcher.h), the tokio-loop replacement for +//! workerd's `--watch` file watcher. The C++ side owns the platform notification fd (inotify on +//! Linux, kqueue on macOS/BSD) and does all the event parsing; the Rust side only supplies +//! "resolve when this fd becomes readable", replacing +//! `kj::UnixEventPort::FdObserver::whenBecomesReadable()`. +//! +//! Semantics: +//! +//! - The fd is registered with the tokio I/O driver per call (edge-triggered underneath, but +//! both epoll and kqueue report readiness that already exists at registration time, so events +//! queued on the fd before the call are not missed). +//! - Dropping the future deregisters the fd without consuming anything. +//! - The caller must keep the fd open until the future resolves or is dropped, and should not +//! have the same fd registered through this function twice concurrently (tokio's I/O driver +//! does not support duplicate registrations of one fd). + +use crate::error::Result; +use crate::runtime::with_runtime; + +/// Resolves when `fd` becomes readable. Unix only; errors immediately on other platforms. +pub async fn wait_fd_readable(fd: i32) -> Result<()> { + #[cfg(unix)] + { + use tokio::io::Interest; + use tokio::io::unix::AsyncFd; + + use crate::error::op; + with_runtime(async move { + let afd = AsyncFd::with_interest(fd, Interest::READABLE).map_err(op("AsyncFd"))?; + // The guard's readiness state is intentionally not cleared: the AsyncFd is + // deregistered immediately below (drop), and the next call re-registers, at which + // point still-pending readiness is reported again. + let _guard = afd.readable().await.map_err(op("readable"))?; + Ok(()) + }) + .await + } + #[cfg(not(unix))] + { + use crate::error::KjIoError; + let _ = fd; + Err(KjIoError::other( + "wait_fd_readable", + "kj-rs-io fd readiness watching is only implemented on Unix", + )) + } +} diff --git a/src/rust/cxx/kj-rs-io/runtime.rs b/src/rust/cxx/kj-rs-io/runtime.rs new file mode 100644 index 00000000000..914a51567ec --- /dev/null +++ b/src/rust/cxx/kj-rs-io/runtime.rs @@ -0,0 +1,50 @@ +//! Runtime-context plumbing: tokio I/O objects must be created (registered with the I/O driver) +//! from within a tokio runtime context, but kj-rs bridge futures are polled by the KJ event +//! loop, outside any `block_on`. These helpers enter this thread's `kj_rs_tokio` runtime context +//! around each poll / each synchronous operation. + +use std::future::Future; + +use tokio::runtime::Handle; + +use crate::error::KjIoError; +use crate::error::Result; + +/// Returns a handle to this thread's KJ-loop tokio runtime, or a `kj::Exception`-convertible +/// error if there is no `TokioEventPort` on this thread. +pub fn runtime_handle() -> Result { + kj_rs_tokio::current_handle().ok_or_else(|| { + KjIoError::other( + "kj_rs_io", + "no kj-rs-tokio runtime on this thread; kj-rs-io requires a TokioEventPort \ + (see kj_rs_io::setupTokioAsyncIo())", + ) + }) +} + +/// Aborts the wrapped tokio task when dropped. Used by the "runtime task forwards a result over +/// a oneshot" pattern (see `net.rs::resolve_host`): if the awaiting bridged future is dropped +/// (KJ promise cancelled), the forwarding task is aborted instead of lingering until its +/// underlying operation completes on its own. +pub struct AbortOnDrop(pub tokio::task::JoinHandle<()>); + +impl Drop for AbortOnDrop { + fn drop(&mut self) { + self.0.abort(); + } +} + +/// Runs `fut` with the current thread's KJ-loop runtime context entered around every poll, so +/// tokio resources created inside it can register with the runtime's I/O driver and timers. +pub async fn with_runtime(fut: impl Future>) -> Result { + let handle = runtime_handle()?; + // Pin the future on the stack, then poll it through `poll_fn` with the runtime guard held + // across each poll. This needs no manual pin-projection (`std::pin::pin!` gives a safe + // `Pin<&mut _>`), so the whole helper is safe. + let mut fut = std::pin::pin!(fut); + std::future::poll_fn(|cx| { + let _guard = handle.enter(); + fut.as_mut().poll(cx) + }) + .await +} diff --git a/src/rust/cxx/kj-rs-io/serve.rs b/src/rust/cxx/kj-rs-io/serve.rs new file mode 100644 index 00000000000..9b1897a3744 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/serve.rs @@ -0,0 +1,431 @@ +//! The native-serve entry points: give a Rust server (any tokio consumer) the +//! best-available tokio-side byte stream for an owned `kj::AsyncIoStream`. +//! +//! Three tiers, two entry points: +//! (1) kj-rs-io-originated streams give up their native tokio socket (the hollow +//! wrapper is destroyed); (2) [`take_kj_socket`] only — a foreign fd-backed stream's fd is +//! duplicated into a fresh tokio socket; (3) [`serve_kj_stream`] only — any other foreign +//! stream is bridged through an in-memory duplex plus a pump future that owns it. Ownership +//! arrives as a [`KjOwn`], so both entry points are safe functions: there is no +//! keep-it-alive caller contract, and the stream is destroyed by the tier that consumed it. +//! See the two entry points' docs for the remaining semantics, in particular why the fd tier +//! is caller-asserted and never automatic. +//! +//! # Pump semantics (matching the hand-built pumps this subsumes) +//! +//! - Bidirectional; each direction ends independently. +//! - Half-close propagates both ways: kj-side EOF shuts down the duplex write half (the tokio +//! consumer reads EOF); the consumer shutting down (or dropping) its duplex end results in +//! `shutdownWrite()` on the kj stream. +//! - Peer-teardown-shaped kj failures (DISCONNECTED reads/writes) are treated as normal EOF, +//! not errors — abrupt client disconnects are normal server load. +//! - Dropping the pump future cancels the in-flight bridged kj promises synchronously, drops +//! the kj-side duplex end (the tokio consumer observes EOF), and destroys the owned kj +//! stream — the peer observes teardown, not a zombie half-open connection (abort-on-drop). + +use std::future::Future; +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; + +use cxx::KjError; +use cxx::KjException; +use cxx::KjExceptionType; +use kj_rs::KjOwn; +use tokio::io::AsyncRead; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWrite; +use tokio::io::AsyncWriteExt; +use tokio::io::DuplexStream; +use tokio::io::ReadBuf; +use tokio::net::TcpStream; +#[cfg(unix)] +use tokio::net::UnixStream; + +use crate::ffi::KjAsyncIoStream; +#[cfg(unix)] +use crate::ffi::dup_raw_fd; +#[cfg(unix)] +use crate::ffi::kj_stream_get_handle; +use crate::ffi::split_kj_stream; +use crate::ffi::unwrap_tokio_stream; + +/// Read chunk size for the pump fallback. +const PUMP_BUF: usize = 8192; + +/// In-memory buffer per direction of the pump's duplex (how far the two sides may run ahead +/// of each other before backpressure). +pub(crate) const DUPLEX_CAPACITY: usize = 4 * PUMP_BUF; + +/// Which transport path [`serve_kj_stream`] produced (perf observability: `Pumped` costs FFI +/// promise round-trips per buffer, `Native` costs none). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ServePath { + /// A native tokio socket (unwrap fast path). + Native, + /// An in-memory duplex fed by the FFI stream pump. + Pumped, +} + +impl std::fmt::Display for ServePath { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::Native => "native", + Self::Pumped => "pumped", + }) + } +} + +/// The tokio-side byte stream for a served kj stream: a native socket or the consumer end of +/// the pump's duplex. +/// +/// Implements `AsyncRead + AsyncWrite`, so it drops into any tokio +/// consumer -- but see [`ServedKjStream::io`] for the thread-affinity contract: only the +/// native variants may be driven off the KJ event-loop thread; `Duplex` is loop-thread-only. +pub enum ServeIo { + Tcp(TcpStream), + #[cfg(unix)] + Unix(UnixStream), + Duplex(DuplexStream), +} + +impl ServeIo { + /// Which transport path this stream is on (see [`ServePath`]). + #[must_use] + pub fn path(&self) -> ServePath { + match self { + Self::Tcp(_) => ServePath::Native, + #[cfg(unix)] + Self::Unix(_) => ServePath::Native, + Self::Duplex(_) => ServePath::Pumped, + } + } + + /// Sets `TCP_NODELAY` on the underlying socket where applicable (a no-op for Unix-domain + /// and duplex transports, which have no Nagle to disable). + /// + /// # Errors + /// + /// Returns the underlying `std::io::Error` if setting the socket option fails. + pub fn set_nodelay(&self, nodelay: bool) -> std::io::Result<()> { + match self { + Self::Tcp(s) => s.set_nodelay(nodelay), + #[cfg(unix)] + Self::Unix(_) => Ok(()), + Self::Duplex(_) => Ok(()), + } + } +} + +impl AsyncRead for ServeIo { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + match self.get_mut() { + Self::Tcp(s) => Pin::new(s).poll_read(cx, buf), + #[cfg(unix)] + Self::Unix(s) => Pin::new(s).poll_read(cx, buf), + Self::Duplex(s) => Pin::new(s).poll_read(cx, buf), + } + } +} + +impl AsyncWrite for ServeIo { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + match self.get_mut() { + Self::Tcp(s) => Pin::new(s).poll_write(cx, buf), + #[cfg(unix)] + Self::Unix(s) => Pin::new(s).poll_write(cx, buf), + Self::Duplex(s) => Pin::new(s).poll_write(cx, buf), + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.get_mut() { + Self::Tcp(s) => Pin::new(s).poll_flush(cx), + #[cfg(unix)] + Self::Unix(s) => Pin::new(s).poll_flush(cx), + Self::Duplex(s) => Pin::new(s).poll_flush(cx), + } + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.get_mut() { + Self::Tcp(s) => Pin::new(s).poll_shutdown(cx), + #[cfg(unix)] + Self::Unix(s) => Pin::new(s).poll_shutdown(cx), + Self::Duplex(s) => Pin::new(s).poll_shutdown(cx), + } + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[std::io::IoSlice<'_>], + ) -> Poll> { + match self.get_mut() { + Self::Tcp(s) => Pin::new(s).poll_write_vectored(cx, bufs), + #[cfg(unix)] + Self::Unix(s) => Pin::new(s).poll_write_vectored(cx, bufs), + Self::Duplex(s) => Pin::new(s).poll_write_vectored(cx, bufs), + } + } + + fn is_write_vectored(&self) -> bool { + match self { + Self::Tcp(s) => s.is_write_vectored(), + #[cfg(unix)] + Self::Unix(s) => s.is_write_vectored(), + Self::Duplex(s) => s.is_write_vectored(), + } + } +} + +/// The KJ-side pump future of the fallback path. +/// +/// Not `Send`: it awaits bridged `kj::Promise`s and must be polled on the KJ event-loop thread +/// owning the stream. Resolves when both directions are done; dropping it aborts the +/// connection bridge (see the module docs). +pub type StreamPump = Pin>>>; + +/// The result of [`serve_kj_stream`]. +pub struct ServedKjStream { + /// The tokio-side stream. + /// + /// Thread affinity (load-bearing, not advisory): the NATIVE variants (`Tcp`/`Unix`) may be + /// handed to a connection task on any runtime -- they stay registered with the I/O driver + /// that created them (for kj-rs-io streams, this thread's KJ-loop runtime) and wake their + /// consumer via tokio's own `Send + Sync` task waker. The [`ServeIo::Duplex`] variant must + /// be driven ONLY on the KJ event-loop thread: its peer end lives inside `pump`, which is + /// polled as a bridged future, so the waker parked in the duplex's internal waker slots is + /// a `kj_rs` `FutureWakerCell` clone -- non-atomic and loop-thread-only by design. A + /// read/write/drop of the duplex from any other thread wakes that cell cross-thread: a data + /// race on its refcount plus a cross-thread `Event::armDepthFirst()` on the KJ event loop + /// (undefined behavior, not merely a logic error). The type is `Send` solely for the native + /// variants' sake; check [`ServedKjStream::path`] before moving it to another thread. + pub io: ServeIo, + /// Present iff `io` is [`ServeIo::Duplex`]: the pump that actually moves the bytes, owning + /// the kj stream it bridges. The caller must poll it on the KJ event-loop thread until it + /// settles or is dropped; dropping it destroys the stream (see the module docs). + pub pump: Option, +} + +impl ServedKjStream { + /// Which transport path was taken (see [`ServePath`]). + #[must_use] + pub fn path(&self) -> ServePath { + self.io.path() + } +} + +// ======================================================================================= +// Entry points + +/// [`take_kj_socket`]'s error: the failure, plus the untouched stream handed back so the +/// caller can fall back to [`serve_kj_stream`]'s pump tier (or destroy it). +pub struct TakeSocketError { + /// The stream `take_kj_socket` consumed, returned untouched. + pub stream: KjOwn, + /// Why the socket could not be taken natively. + pub error: KjError, +} + +impl std::fmt::Debug for TakeSocketError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TakeSocketError") + .field("error", &self.error) + .finish_non_exhaustive() + } +} + +impl std::fmt::Display for TakeSocketError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self.error) + } +} + +/// Drops the handed-back stream (on the current — KJ event-loop — thread) and keeps the error. +impl From for KjError { + fn from(e: TakeSocketError) -> Self { + e.error + } +} + +/// Tier-1 unwrap: if the owned stream is kj-rs-io-originated, moves its native tokio socket +/// out (leaving the C++ wrapper hollow) and returns it. `None` for a foreign stream. +fn unwrap_native(stream: &mut KjOwn) -> Option { + unwrap_tokio_stream(stream.as_mut()) + .ok() + .and_then(|native| native.into_serve_io()) +} + +/// Takes the stream's socket natively (tiers 1 + 2). +/// +/// Tier 1 moves the native tokio object out of a kj-rs-io stream; tier 2 — unix only — +/// duplicates the stream's OS fd (`F_DUPFD_CLOEXEC`, forced non-blocking) into a fresh tokio +/// socket. Either way the result is an owned, KJ-independent [`ServeIo`] (never +/// [`ServeIo::Duplex`]) and the consumed kj stream is destroyed before returning. +/// +/// The handle tier (tier 2) is `cfg(unix)`: under the all-rust mode on Windows every +/// socket-backed stream originates in kj-rs-io, so tier 1 always applies and a windows dup arm +/// would be dead code (see `dup_raw_fd`'s docs for the analysis and the +/// `BorrowedSocket::try_clone_to_owned` escape hatch should that ever change). +/// +/// The caller asserts the stream is a **plain stream socket**: if it exposes an fd, that fd +/// carries the stream's own bytes. Byte-transforming wrappers (TLS — `kj::TlsConnection` +/// forwards `getFd()` to the ciphertext transport socket) violate this and must go through +/// [`serve_kj_stream`] instead. As with destroying any kj stream, no I/O promises may be +/// outstanding on it when ownership is handed over. +/// +/// # Errors +/// +/// Errors when the stream is neither kj-rs-io native nor fd-backed (in-memory pipes, promised +/// streams, non-Unix platforms' foreign streams): such transports can only be served through +/// [`serve_kj_stream`]'s pump tier — the error hands the stream back for exactly that +/// fallback. Also errors if the dup or tokio registration fails. +pub fn take_kj_socket( + stream: KjOwn, +) -> std::result::Result { + let mut stream = stream; + if let Some(io) = unwrap_native(&mut stream) { + // Tier 1: the wrapper is hollow; destroy it now. + drop(stream); + return Ok(io); + } + #[cfg(unix)] + { + let handle = kj_stream_get_handle(&stream); + if handle >= 0 { + // A unix fd is a non-negative int, widened losslessly to i64 by the bridge. + #[expect(clippy::cast_possible_truncation)] + let fd = handle as i32; + let result = dup_raw_fd(fd) + .map_err(KjError::from) + .and_then(|owned| crate::net::serve_io_from_owned_fd(owned).map_err(KjError::from)); + return match result { + Ok(io) => { + // Tier 2: the dup is independent; the original stream (and its fd) can go. + drop(stream); + Ok(io) + } + Err(error) => Err(TakeSocketError { stream, error }), + }; + } + } + Err(TakeSocketError { + stream, + error: KjError::new( + KjExceptionType::Failed, + "cannot take the stream's socket natively: not a kj-rs-io stream and no underlying \ + OS fd (foreign non-socket transport; serve it through serve_kj_stream's pump instead)" + .to_owned(), + ), + }) +} + +/// Yields the best-available tokio-side stream for the owned `stream`. +/// +/// That is the native tokio object when `stream` originated in kj-rs-io, else an in-memory +/// duplex bridged by a pump future that owns the stream. Never extracts a foreign stream's fd +/// — kj wrappers forward `getFd()` to their transport socket, so a byte-transforming wrapper's +/// fd carries the wrong bytes (TLS ciphertext); callers that can assert a plain socket should +/// prefer [`take_kj_socket`]. +/// +/// On the native path the hollow wrapper is destroyed before returning; on the pump path the +/// stream lives inside the pump and is destroyed when the pump settles or is dropped. The pump +/// must only be polled from the KJ event-loop thread owning the stream. As with destroying any +/// kj stream, no I/O promises may be outstanding on it when ownership is handed over. +#[must_use] +pub fn serve_kj_stream(stream: KjOwn) -> ServedKjStream { + let mut stream = stream; + if let Some(io) = unwrap_native(&mut stream) { + // Native path: the wrapper is hollow; destroy it now. + drop(stream); + return ServedKjStream { io, pump: None }; + } + + // Foreign stream (or, pathologically, an already-hollow wrapper, whose pump reads will + // surface the "already unwrapped" error): bridge through a duplex pump owning the stream. + let (consumer_end, kj_end) = tokio::io::duplex(DUPLEX_CAPACITY); + let pump = Box::pin(pump_kj_stream(stream, kj_end)); + ServedKjStream { + io: ServeIo::Duplex(consumer_end), + pump: Some(pump), + } +} + +/// Whether a bridged kj exception is peer-teardown-shaped (treated as EOF by the pump). +fn is_disconnected(exception: &KjException) -> bool { + exception.r#type() == KjExceptionType::Disconnected +} + +/// The duplex pump: bridges the owned `stream` (via bridged `kj::io` promises, on the calling +/// KJ thread) to `kj_end`, the pump-side end of the consumer's duplex. Owns the stream: it is +/// destroyed when this future settles or is dropped. +/// +/// Unsafe-free and compiler-checked end to end: the stream is split into typed read/write +/// halves ([`split_kj_stream`]), so the borrow checker enforces kj's stream contract — at most +/// one read and one write in flight, nothing else touching the stream while the halves live — +/// and proves the owner outlives every in-flight bridged promise. +pub(crate) async fn pump_kj_stream( + mut stream: KjOwn, + kj_end: DuplexStream, +) -> Result<(), KjError> { + let (mut rd, mut wr) = split_kj_stream(&mut stream); + let (mut from_consumer, mut to_consumer) = tokio::io::split(kj_end); + + // kj stream -> consumer. Ends (shutting down the duplex write half, i.e. EOF to the + // consumer) at kj-side EOF — or when the peer disconnects abruptly (DISCONNECTED read + // failures are normal client behavior, treated as EOF). + let kj_to_consumer = async { + let mut buf = vec![0u8; PUMP_BUF]; + loop { + let n = match rd.try_read(&mut buf, 1).await { + Ok(n) => n, + Err(e) if is_disconnected(&e) => 0, + Err(e) => return Err(KjError::from(e)), + }; + if n == 0 { + let _ = to_consumer.shutdown().await; + return Ok::<(), KjError>(()); + } + if to_consumer.write_all(&buf[..n]).await.is_err() { + // The consumer dropped its duplex end: the connection was abandoned + // deliberately; nothing more to deliver in this direction. + return Ok(()); + } + } + }; + + // consumer -> kj stream. Ends (with a kj-side shutdownWrite) when the consumer shuts + // down or drops its end — or, without an error, when the kj peer already went away (a + // DISCONNECTED write failure: the consumer's remaining output has nowhere to go). + let consumer_to_kj = async { + let mut buf = vec![0u8; PUMP_BUF]; + loop { + let n = match from_consumer.read(&mut buf).await { + // Duplex reads only fail if the consumer end vanished ungracefully; either + // way this direction is over. + Ok(0) | Err(_) => 0, + Ok(n) => n, + }; + if n == 0 { + wr.shutdown_write(); + return Ok::<(), KjError>(()); + } + match wr.write(&buf[..n]).await { + Ok(()) => {} + Err(e) if is_disconnected(&e) => return Ok(()), + Err(e) => return Err(KjError::from(e)), + } + } + }; + + tokio::try_join!(kj_to_consumer, consumer_to_kj).map(|((), ())| ()) +} diff --git a/src/rust/cxx/kj-rs-io/signal.rs b/src/rust/cxx/kj-rs-io/signal.rs new file mode 100644 index 00000000000..f4169dccaa7 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/signal.rs @@ -0,0 +1,144 @@ +//! tokio-backed signal watching: POSIX signals on Unix, the corresponding console control +//! events on Windows. +//! +//! This backs `kj_rs_io::onSignal()` (async-io.h), the tokio-loop replacement for +//! `kj::UnixEventPort::onSignal()` -- workerd uses it for SIGTERM graceful drain. +//! +//! Semantics differences vs `UnixEventPort::onSignal()` (acceptable for the drain use case): +//! +//! - No `siginfo_t` is reported; the promise just resolves. +//! - The handler is registered when the returned future is first polled (tokio registers with +//! the process-global signal registry at `signal()` time), not at call time, and KJ does not +//! block the signal beforehand the way `UnixEventPort::captureSignal()` does. A signal +//! delivered before the first poll takes its default disposition. +//! - tokio's signal registration is process-wide and persists for the life of the process +//! (dropping the future stops *watching*, but does not restore `SIG_DFL`). +//! +//! Both arms route the `recv()` through a tokio runtime task rather than awaiting the stream +//! from the bridged future: tokio's signal registry is process-global and its broadcast can run +//! on a different thread (another runtime's loop thread on unix, the console-ctrl thread on +//! Windows), which must never wake a bridged loop-thread-only waker directly. See the comments +//! in each arm. +//! +//! On Windows the signums workerd actually passes are mapped to their conventional console +//! control events: SIGTERM -> `ctrl_shutdown`, SIGINT -> `ctrl_c`. Anything else errors. + +use crate::error::KjIoError; +use crate::error::Result; +use crate::runtime::with_runtime; + +/// Resolves when the process receives signal `signum` (on Windows: the console control event +/// conventionally mapped to it). Errors immediately for unmapped signums / other platforms. +pub async fn wait_for_signal(signum: i32) -> Result<()> { + #[cfg(unix)] + { + use crate::error::op; + with_runtime(async move { + let kind = tokio::signal::unix::SignalKind::from_raw(signum); + // Create the stream here (inside the runtime context, at first poll of the bridged + // future) so the process-global handler registration happens as early as possible. + let mut sig = tokio::signal::unix::signal(kind).map_err(op("signal"))?; + + // Do NOT `sig.recv().await` directly from this (bridged) future. tokio's signal + // registry is process-global: when several tokio runtimes exist in the process (one + // per KJ event loop thread — e.g. workerd's main loop plus the inspector thread's + // loop), the runtime whose driver consumes the signal's wake byte performs the + // broadcast, so the stored waker can be woken FROM THAT OTHER THREAD. A directly + // parked waker here would be the bridged future's loop-thread-only, non-atomic + // `kj_rs` `FutureWakerCell`: waking it cross-thread is UB under the bridge's + // single-thread waker axiom, and in practice loses the wakeup (observed as workerd + // ignoring SIGTERM whenever the inspector thread's runtime won the race). So, like + // `net.rs::resolve_host` (the model citizen for this pattern), a tokio *runtime* + // task owns the `recv()`: the cross-thread broadcast terminates at tokio's own + // `Send + Sync` scheduler waker (which unparks this loop), the task then runs on the + // loop thread and hands the result back over a oneshot, waking the bridged future + // same-thread. + let (tx, rx) = tokio::sync::oneshot::channel::>(); + let task = tokio::spawn(async move { + let result = sig + .recv() + .await + .ok_or_else(|| KjIoError::other("signal", "signal stream closed unexpectedly")); + let _ = tx.send(result); + }); + // If this future is dropped (KJ promise cancelled), abort the watcher task so its + // signal-stream registration is torn down instead of lingering for the process + // lifetime. + let _abort_guard = crate::runtime::AbortOnDrop(task); + match rx.await { + Ok(result) => result, + Err(_) => Err(KjIoError::other("signal", "signal watcher task dropped")), + } + }) + .await + } + // Validated by Windows CI. + // + // Same forwarding-task pattern as the unix arm, for the Windows flavor of the same hazard: + // tokio's `SetConsoleCtrlHandler` handler runs on an OS-spawned console-ctrl thread and + // broadcasts to every registered watcher's stored waker FROM THAT THREAD. Awaiting the + // stream directly here would park a clone of the bridged future's waker -- a + // loop-thread-only, non-atomic `kj_rs` `FutureWakerCell` -- in tokio's signal registry, and + // one Ctrl-C/shutdown event would wake it cross-thread: UB under the bridge's single-thread + // waker axiom. The tokio *runtime* task owning the `recv()` terminates the cross-thread + // broadcast at tokio's own `Send + Sync` scheduler waker (which unparks this loop); the task + // then runs on the loop thread and hands the result back over a oneshot, waking the bridged + // future same-thread. + #[cfg(windows)] + { + use crate::error::op; + // `` values as the C++ callers pass them (MSVC defines SIGINT=2, SIGTERM=15). + // workerd's only caller passes SIGTERM (graceful drain; server/cli-io-backend.c++); + // SIGINT is mapped for completeness. + const SIGINT: i32 = 2; + const SIGTERM: i32 = 15; + with_runtime(async move { + if signum != SIGTERM && signum != SIGINT { + return Err(KjIoError::other( + "signal", + "kj-rs-io only watches SIGTERM/SIGINT on Windows", + )); + } + let (tx, rx) = tokio::sync::oneshot::channel::>(); + let task = tokio::spawn(async move { + let result = async { + // SIGTERM -> ctrl_shutdown, SIGINT -> ctrl_c (the conventional mappings). + let received = match signum { + SIGTERM => { + let mut sig = + tokio::signal::windows::ctrl_shutdown().map_err(op("signal"))?; + sig.recv().await + } + // SIGINT; anything else already errored before the spawn. + _ => { + let mut sig = tokio::signal::windows::ctrl_c().map_err(op("signal"))?; + sig.recv().await + } + }; + received.ok_or_else(|| { + KjIoError::other("signal", "signal stream closed unexpectedly") + }) + } + .await; + let _ = tx.send(result); + }); + // If this future is dropped (KJ promise cancelled), abort the watcher task so its + // signal-stream registration is torn down instead of lingering for the process + // lifetime. + let _abort_guard = crate::runtime::AbortOnDrop(task); + match rx.await { + Ok(result) => result, + Err(_) => Err(KjIoError::other("signal", "signal watcher task dropped")), + } + }) + .await + } + #[cfg(not(any(unix, windows)))] + { + let _ = signum; + Err(KjIoError::other( + "signal", + "kj-rs-io signal watching is not implemented on this platform", + )) + } +} diff --git a/src/rust/cxx/kj-rs-io/stream.rs b/src/rust/cxx/kj-rs-io/stream.rs new file mode 100644 index 00000000000..77aa99b30ba --- /dev/null +++ b/src/rust/cxx/kj-rs-io/stream.rs @@ -0,0 +1,559 @@ +//! Tokio-backed byte streams behind KJ's stream interfaces. +//! +//! `TokioStream` (TCP or Unix domain) implements the `kj::AsyncIoStream` operations; all I/O +//! uses tokio's `&self` readiness API (`ready()` + `try_read`/`try_write`), which supports +//! concurrent reads and writes on one stream and is cancel-safe: dropping a pending future +//! (i.e. dropping the wrapping `kj::Promise`) merely deregisters the waker, releasing the +//! readiness interest so the stream can be reused or dropped sanely. + +use std::io::Read; +use std::io::Write; + +use tokio::io::Interest; +use tokio::net::TcpStream; +#[cfg(unix)] +use tokio::net::UnixStream; + +use crate::error::KjIoError; +use crate::error::Result; +use crate::error::op; +use crate::runtime::with_runtime; + +/// A native tokio stream (the "unwrap fast path" object). +/// +/// C++ holds one of these inside every kj-rs-io `kj::AsyncIoStream`; Rust code can take it +/// back out via [`crate::unwrap_kj_stream`] and drive the connection natively. +pub struct TokioStream { + /// `None` after the native stream has been moved out by [`TokioStream::take`] (the C++ + /// wrapper is then "hollow" and every operation fails). + inner: Option, +} + +enum Inner { + Tcp(TcpStream), + #[cfg(unix)] + Unix(UnixStream), +} + +impl TokioStream { + #[must_use] + pub fn from_tcp(stream: TcpStream) -> Self { + Self { + inner: Some(Inner::Tcp(stream)), + } + } + + #[cfg(unix)] + #[must_use] + pub fn from_unix(stream: UnixStream) -> Self { + Self { + inner: Some(Inner::Unix(stream)), + } + } + + /// Recovers the native tokio TCP stream, if this is a (non-hollow) TCP stream. + #[must_use] + pub fn into_tcp_stream(self) -> Option { + match self.inner { + Some(Inner::Tcp(stream)) => Some(stream), + _ => None, + } + } + + /// Recovers the native tokio Unix-domain stream, if this is one. + #[cfg(unix)] + #[must_use] + pub fn into_unix_stream(self) -> Option { + match self.inner { + Some(Inner::Unix(stream)) => Some(stream), + _ => None, + } + } + + fn inner(&self) -> Result<&Inner> { + self.inner + .as_ref() + .ok_or_else(|| KjIoError::other("kj_rs_io", "stream was unwrapped (hollow wrapper)")) + } + + async fn ready(&self, interest: Interest) -> Result<()> { + match self.inner()? { + Inner::Tcp(s) => s.ready(interest).await, + #[cfg(unix)] + Inner::Unix(s) => s.ready(interest).await, + } + .map_err(op("poll()"))?; + Ok(()) + } + + #[expect( + clippy::expect_used, + reason = "only called from try_read_min/write_all, which call self.inner()? first, so `inner` is Some here; None occurs only for a hollow (unwrapped) wrapper, which is never read" + )] + fn try_read(&self, buf: &mut [u8]) -> std::io::Result { + match self.inner.as_ref().expect("checked by caller") { + Inner::Tcp(s) => s.try_read(buf), + #[cfg(unix)] + Inner::Unix(s) => s.try_read(buf), + } + } + + #[expect( + clippy::expect_used, + reason = "only called from write_all, which calls self.inner()? first, so `inner` is Some here; None occurs only for a hollow (unwrapped) wrapper, which is never written" + )] + fn try_write(&self, buf: &[u8]) -> std::io::Result { + match self.inner.as_ref().expect("checked by caller") { + Inner::Tcp(s) => s.try_write(buf), + #[cfg(unix)] + Inner::Unix(s) => s.try_write(buf), + } + } + + /// KJ `tryRead` semantics: loop until at least `min_bytes` (or EOF), up to `buf.len()`. + async fn try_read_min(&self, buf: &mut [u8], min_bytes: usize) -> Result { + self.inner()?; + let min_bytes = min_bytes.min(buf.len()); + let mut total = 0; + while total < min_bytes { + match self.try_read(&mut buf[total..]) { + Ok(0) => break, // EOF: return what we have (< min_bytes signals EOF to KJ). + Ok(n) => total += n, + Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => { + self.ready(Interest::READABLE).await?; + } + Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {} + Err(e) => return Err(op("read()")(e)), + } + } + Ok(total) + } + + /// Write-all semantics. + async fn write_all(&self, buf: &[u8]) -> Result<()> { + self.inner()?; + let mut written = 0; + while written < buf.len() { + match self.try_write(&buf[written..]) { + Ok(0) => { + // try_write on a socket signals "would block" via Err(WouldBlock), so a + // zero-byte result for a non-empty buffer means the connection is gone. + return Err(KjIoError::other( + "write()", + "wrote zero bytes (connection closed)", + )); + } + Ok(n) => written += n, + Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => { + self.ready(Interest::WRITABLE).await?; + } + Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {} + Err(e) => return Err(op("write()")(e)), + } + } + Ok(()) + } + + /// Resolves once new writes are doomed to fail (peer reset / hangup observed). + /// + /// tokio has no direct primitive for this, so we register a *duplicate* of the socket fd + /// with the I/O driver for WRITABLE interest and wait — explicitly clearing plain-writable + /// readiness — until the OS reports write-closed (kqueue: `EV_EOF` on the write filter; + /// epoll: `EPOLLHUP`/`EPOLLERR`) or an error. Like KJ's own implementation this does *not* + /// fire on a mere half-close (peer FIN / `EPOLLRDHUP`): reads hitting EOF must not count as + /// write-disconnect. + /// + /// Windows behavior (see the arm below): a never-resolving future, which IS KJ + /// parity — KJ's *current* Windows behavior (`whenWriteDisconnected` returns `NEVER_DONE`). + /// Win32 has no documented primitive for detecting disconnect without a read/write; KJ's own + /// TODO points at the undocumented-but-stable `IOCTL_AFD_POLL` ioctl (capnproto + /// `async-io-win32.c++:289` — the mechanism `select()` itself is built on). An AFD-poll + /// implementation remains an optional upgrade over this documented parity behavior. + #[cfg(unix)] + async fn when_write_disconnected(&self) -> Result<()> { + use std::os::fd::AsRawFd; + + // Borrow the live socket's fd directly (tokio streams implement `AsFd`) and dup it, so + // no raw fd is ever materialized without an owner. + let borrowed = self.as_borrowed_fd()?; + let owned = borrowed.try_clone_to_owned().map_err(op("dup()"))?; + debug_assert_ne!(owned.as_raw_fd(), borrowed.as_raw_fd()); + // The dup shares the underlying open socket (and its O_NONBLOCK status), but has its own + // registration with the I/O driver, so clearing readiness here never disturbs reads or + // writes on the primary registration. + let async_fd = tokio::io::unix::AsyncFd::with_interest(owned, Interest::WRITABLE) + .map_err(op("whenWriteDisconnected"))?; + loop { + let mut guard = async_fd + .ready(Interest::WRITABLE) + .await + .map_err(op("whenWriteDisconnected"))?; + let ready = guard.ready(); + if ready.is_write_closed() || ready.is_error() { + return Ok(()); + } + // Plain "writable": clear it so the next wait sleeps until an actual state-change + // event (edge-triggered), rather than spinning on an always-writable socket. + guard.clear_ready(); + } + } + + /// Never resolves: KJ parity, not a gap — capnproto's win32 `whenWriteDisconnected` returns + /// `NEVER_DONE` today (its `IOCTL_AFD_POLL` idea is only a TODO; see the Windows-behavior note + /// on the unix arm above). Validated by Windows CI. + #[cfg(windows)] + async fn when_write_disconnected(&self) -> Result<()> { + self.inner()?; + std::future::pending::<()>().await; + unreachable!() + } + + #[cfg(not(any(unix, windows)))] + async fn when_write_disconnected(&self) -> Result<()> { + self.inner()?; + Err(KjIoError::other( + "whenWriteDisconnected", + "not implemented by kj-rs-io on this platform", + )) + } + + fn shutdown_write(&self) -> Result<()> { + #[cfg(unix)] + { + let dup = self + .as_borrowed_fd()? + .try_clone_to_owned() + .map_err(op("dup()"))?; + // shutdown() acts on the socket itself, so performing it through a dup'd fd + // affects the shared socket, and dropping the dup only closes the duplicate. + let result = match self.inner()? { + Inner::Tcp(_) => std::net::TcpStream::from(dup).shutdown(std::net::Shutdown::Write), + Inner::Unix(_) => { + std::os::unix::net::UnixStream::from(dup).shutdown(std::net::Shutdown::Write) + } + }; + result.map_err(op("shutdown(SHUT_WR)")) + } + // Validated by Windows CI; mirrors the unix arm. No dup: `with_sock_ref` borrows the + // live socket (`SockRef`), and winsock `shutdown` acts on the underlying socket either + // way — the unix arm dups only because std's `shutdown` is a method on owning types, + // whereas the windows `BorrowedSocket::try_clone_to_owned` equivalent + // (WSADuplicateSocketW) would be strictly heavier than the borrow. + #[cfg(windows)] + { + self.with_sock_ref("shutdown(SD_SEND)", |sock| { + sock.shutdown(std::net::Shutdown::Write) + }) + } + #[cfg(not(any(unix, windows)))] + { + Err(KjIoError::other( + "shutdownWrite", + "not implemented on this platform", + )) + } + } + + /// Borrows the live tokio socket's fd (tokio streams implement `AsFd`), for dup-based + /// operations that must not conjure a raw fd out of an integer. Errs if the wrapper is + /// hollow. + #[cfg(unix)] + pub(crate) fn as_borrowed_fd(&self) -> Result> { + use std::os::fd::AsFd; + Ok(match self.inner()? { + Inner::Tcp(s) => s.as_fd(), + Inner::Unix(s) => s.as_fd(), + }) + } + + /// Borrows the live tokio socket's `SOCKET` (tokio's `TcpStream` implements `AsSocket`): + /// the Windows counterpart of [`TokioStream::as_borrowed_fd`]. Errs if the wrapper is + /// hollow. On Windows only the Tcp variant of `Inner` exists. + // Validated by Windows CI; mirrors the unix arm. + #[cfg(windows)] + pub(crate) fn as_borrowed_socket(&self) -> Result> { + use std::os::windows::io::AsSocket; + Ok(match self.inner()? { + Inner::Tcp(s) => s.as_socket(), + }) + } + + /// Runs `f` on a `socket2::SockRef` borrowing the live socket — the shared body of the + /// `getsockname()`/`getpeername()` passthroughs; only the socket borrow is per-platform. + fn with_sock_ref( + &self, + op_name: &'static str, + f: impl FnOnce(&socket2::SockRef<'_>) -> std::io::Result, + ) -> Result { + #[cfg(unix)] + { + let fd = self.as_borrowed_fd()?; + f(&socket2::SockRef::from(&fd)).map_err(op(op_name)) + } + // Validated by Windows CI; mirrors the unix arm. + #[cfg(windows)] + { + let sock = self.as_borrowed_socket()?; + f(&socket2::SockRef::from(&sock)).map_err(op(op_name)) + } + #[cfg(not(any(unix, windows)))] + { + let _ = f; + self.inner()?; + Err(KjIoError::other( + op_name, + "not implemented by kj-rs-io on this platform", + )) + } + } + + /// Raw `struct sockaddr` bytes of the socket's locally-bound address (the `getsockname()` + /// passthrough behind `kj::AsyncIoStream::getsockname`). + fn local_addr_bytes(&self) -> Result> { + let addr = self.with_sock_ref("getsockname()", |sock| sock.local_addr())?; + Ok(crate::ffi::sockaddr_to_bytes(&addr)) + } + + /// Raw `struct sockaddr` bytes of the connected peer's address (the `getpeername()` + /// passthrough behind `kj::AsyncIoStream::getpeername` and the accept-loop peer-filter + /// check). + fn peer_addr_bytes(&self) -> Result> { + let addr = self.with_sock_ref("getpeername()", |sock| sock.peer_addr())?; + Ok(crate::ffi::sockaddr_to_bytes(&addr)) + } + + /// The underlying raw OS socket handle, widened to `i64`: a Unix fd + /// (`kj::AsyncIoStream::getFd()`) or a win32 `SOCKET` (`getWin32Handle()`). + fn raw_handle(&self) -> Result { + #[cfg(unix)] + { + use std::os::fd::AsRawFd; + Ok(i64::from(match self.inner()? { + Inner::Tcp(s) => s.as_raw_fd(), + Inner::Unix(s) => s.as_raw_fd(), + })) + } + // Validated by Windows CI; mirrors the unix arm. + #[cfg(windows)] + { + use std::os::windows::io::AsRawSocket; + let raw = match self.inner()? { + Inner::Tcp(s) => s.as_raw_socket(), + }; + // A live SOCKET fits in i64 (Windows handles fit in 32 bits); the bridge carries + // its bits verbatim. + #[allow(clippy::cast_possible_wrap)] + Ok(raw as i64) + } + #[cfg(not(any(unix, windows)))] + { + self.inner()?; + Err(KjIoError::other( + "getFd", + "file descriptors are not available on this platform", + )) + } + } + + /// Recovers whichever native tokio object this is, as a [`crate::serve::ServeIo`] + /// (the unwrap fast path of [`crate::serve_kj_stream`]). `None` if hollow. + pub(crate) fn into_serve_io(self) -> Option { + match self.inner? { + Inner::Tcp(stream) => Some(crate::serve::ServeIo::Tcp(stream)), + #[cfg(unix)] + Inner::Unix(stream) => Some(crate::serve::ServeIo::Unix(stream)), + } + } + + fn take(&mut self) -> Result> { + let inner = self.inner.take().ok_or_else(|| { + KjIoError::other("kj_rs_io", "stream was already unwrapped (hollow wrapper)") + })?; + Ok(Box::new(Self { inner: Some(inner) })) + } +} + +// ====================================================================================== +// Bridge entry points (see lib.rs). Every async fn wraps its body in `with_runtime` so tokio +// resources created while polling on the KJ thread can reach the loop runtime's I/O driver. + +pub async fn stream_try_read( + stream: &TokioStream, + buf: &mut [u8], + min_bytes: usize, +) -> Result { + with_runtime(stream.try_read_min(buf, min_bytes)).await +} + +pub async fn stream_write(stream: &TokioStream, buf: &[u8]) -> Result<()> { + with_runtime(stream.write_all(buf)).await +} + +pub async fn stream_when_write_disconnected(stream: &TokioStream) -> Result<()> { + with_runtime(stream.when_write_disconnected()).await +} + +pub fn stream_shutdown_write(stream: &TokioStream) -> Result<()> { + stream.shutdown_write() +} + +pub fn stream_raw_handle(stream: &TokioStream) -> Result { + stream.raw_handle() +} + +pub fn stream_local_addr(stream: &TokioStream) -> Result> { + stream.local_addr_bytes() +} + +pub fn stream_peer_addr(stream: &TokioStream) -> Result> { + stream.peer_addr_bytes() +} + +pub fn stream_take(stream: &mut TokioStream) -> Result> { + stream.take() +} + +// ====================================================================================== +// Arbitrary readable/writable fds (kj::LowLevelAsyncIoProvider::wrapInputFd/wrapOutputFd). +// Unix only: implemented over AsyncFd, which supports pipes, character devices and sockets +// (regular files are rejected by epoll/kqueue, matching KJ's fd-observer-based provider). +// Deliberately no windows arm: kj's win32 LowLevelAsyncIoProvider has no pipe-fd tier — its +// `Fd` is documented as a SOCKET (capnproto async-io.h) and its wrapInputFd/wrapOutputFd are +// implemented identically to wrapSocketFd (async-io-win32.c++) — so the C++ side +// (async-io.c++) routes win32 wrapInputFd/wrapOutputFd through the tested socket path +// (`wrap_socket_fd`) and never calls these entry points there; the `not(unix)` arms below are +// totality backstops only. + +#[cfg(unix)] +type FdIo = tokio::io::unix::AsyncFd; + +pub struct TokioInputFd { + #[cfg(unix)] + inner: FdIo, +} + +pub struct TokioOutputFd { + #[cfg(unix)] + inner: FdIo, +} + +#[cfg(unix)] +fn fd_io_from_raw(fd: i32, interest: Interest) -> Result { + let owned = crate::ffi::own_fd_from_raw(fd); + let _guard = crate::runtime::runtime_handle()?.enter(); + tokio::io::unix::AsyncFd::with_interest(std::fs::File::from(owned), interest) + .map_err(op("wrapFd")) +} + +pub fn wrap_input_fd(fd: i32) -> Result> { + #[cfg(unix)] + { + Ok(Box::new(TokioInputFd { + inner: fd_io_from_raw(fd, Interest::READABLE)?, + })) + } + #[cfg(not(unix))] + { + let _ = fd; + Err(KjIoError::other( + "wrapInputFd", + "not implemented on this platform", + )) + } +} + +pub fn wrap_output_fd(fd: i32) -> Result> { + #[cfg(unix)] + { + Ok(Box::new(TokioOutputFd { + inner: fd_io_from_raw(fd, Interest::WRITABLE)?, + })) + } + #[cfg(not(unix))] + { + let _ = fd; + Err(KjIoError::other( + "wrapOutputFd", + "not implemented on this platform", + )) + } +} + +pub async fn input_fd_try_read( + stream: &TokioInputFd, + buf: &mut [u8], + min_bytes: usize, +) -> Result { + #[cfg(unix)] + { + with_runtime(async move { + let min_bytes = min_bytes.min(buf.len()); + let mut total = 0; + while total < min_bytes { + let mut guard = stream + .inner + .ready(Interest::READABLE) + .await + .map_err(op("poll()"))?; + match guard.try_io(|inner| { + let mut file: &std::fs::File = inner.get_ref(); + file.read(&mut buf[total..]) + }) { + Ok(Ok(0)) => break, // EOF + Ok(Ok(n)) => total += n, + Ok(Err(e)) if e.kind() == std::io::ErrorKind::Interrupted => {} + Ok(Err(e)) => return Err(op("read()")(e)), + Err(_would_block) => {} + } + } + Ok(total) + }) + .await + } + #[cfg(not(unix))] + { + let _ = (stream, buf, min_bytes); + Err(KjIoError::other( + "read()", + "not implemented on this platform", + )) + } +} + +pub async fn output_fd_write(stream: &TokioOutputFd, buf: &[u8]) -> Result<()> { + #[cfg(unix)] + { + with_runtime(async move { + let mut written = 0; + while written < buf.len() { + let mut guard = stream + .inner + .ready(Interest::WRITABLE) + .await + .map_err(op("poll()"))?; + match guard.try_io(|inner| { + let mut file: &std::fs::File = inner.get_ref(); + file.write(&buf[written..]) + }) { + Ok(Ok(0)) => { + return Err(KjIoError::other("write()", "wrote zero bytes")); + } + Ok(Ok(n)) => written += n, + Ok(Err(e)) if e.kind() == std::io::ErrorKind::Interrupted => {} + Ok(Err(e)) => return Err(op("write()")(e)), + Err(_would_block) => {} + } + } + Ok(()) + }) + .await + } + #[cfg(not(unix))] + { + let _ = (stream, buf); + Err(KjIoError::other( + "write()", + "not implemented on this platform", + )) + } +} diff --git a/src/rust/cxx/kj-rs-io/tests/BUILD.bazel b/src/rust/cxx/kj-rs-io/tests/BUILD.bazel new file mode 100644 index 00000000000..a12f68d05ce --- /dev/null +++ b/src/rust/cxx/kj-rs-io/tests/BUILD.bazel @@ -0,0 +1,138 @@ +load("@rules_cc//cc:cc_test.bzl", "cc_test") +load("@rules_rust//rust:defs.bzl", "rust_library") +load("//src/rust/cxx/tools/bazel:rust_cxx_bridge.bzl", "rust_cxx_bridge") + +rust_library( + name = "tests", + srcs = glob(["*.rs"]), + edition = "2024", + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + deps = [ + ":bridge", + "//src/rust/cxx", + "//src/rust/cxx/kj-rs", + "//src/rust/cxx/kj-rs-io", + "//src/rust/cxx/kj-rs-tokio", + "@crates_vendor//:bytes", + "@crates_vendor//:tokio", + ], +) + +rust_cxx_bridge( + name = "bridge", + src = "lib.rs", + include_prefix = "kj-rs-io-test", + deps = [ + "//src/rust/cxx/kj-rs", + "//src/rust/cxx/kj-rs-io:bridge", + ], +) + +cc_test( + name = "async-io-test", + # medium: expectConnectFailure carries a 30 s diagnostic bound and a 60 s watchdog for a + # Windows CI wedge; small's 60 s budget would race the watchdog. + size = "medium", + srcs = ["async-io-test.c++"], + linkstatic = select({ + "@platforms//os:windows": True, + "//conditions:default": False, + }), + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + deps = [ + ":bridge", + ":tests", + "//src/rust/cxx/kj-rs-io:kj-rs-io-lib", + "//src/rust/cxx/third-party:runtime", + "@capnp-cpp//src/kj:kj-test", + ], +) + +cc_test( + name = "file-watcher-test", + size = "small", + srcs = ["file-watcher-test.c++"], + linkstatic = select({ + "@platforms//os:windows": True, + "//conditions:default": False, + }), + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + deps = [ + ":tests", + "//src/rust/cxx/kj-rs-io:kj-rs-io-lib", + "//src/rust/cxx/third-party:runtime", + "@capnp-cpp//src/kj:kj-test", + ], +) + +cc_test( + name = "http-test", + size = "small", + srcs = ["http-test.c++"], + linkstatic = select({ + "@platforms//os:windows": True, + "//conditions:default": False, + }), + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + deps = [ + ":tests", + "//src/rust/cxx/kj-rs-io:kj-rs-io-lib", + "//src/rust/cxx/third-party:runtime", + "@capnp-cpp//src/kj:kj-test", + "@capnp-cpp//src/kj/compat:kj-http", + ], +) + +cc_test( + name = "serve-test", + size = "small", + srcs = ["serve-test.c++"], + linkstatic = select({ + "@platforms//os:windows": True, + "//conditions:default": False, + }), + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + deps = [ + ":bridge", + ":tests", + "//src/rust/cxx/kj-rs-io:kj-rs-io-lib", + "//src/rust/cxx/third-party:runtime", + "@capnp-cpp//src/kj:kj-test", + ], +) + +cc_test( + name = "capnp-rpc-test", + size = "small", + srcs = ["capnp-rpc-test.c++"], + linkstatic = select({ + "@platforms//os:windows": True, + "//conditions:default": False, + }), + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + deps = [ + ":tests", + "//src/rust/cxx/kj-rs-io:kj-rs-io-lib", + "//src/rust/cxx/third-party:runtime", + "@capnp-cpp//src/capnp:capnp-rpc", + "@capnp-cpp//src/kj:kj-test", + ], +) diff --git a/src/rust/cxx/kj-rs-io/tests/async-io-test.c++ b/src/rust/cxx/kj-rs-io/tests/async-io-test.c++ new file mode 100644 index 00000000000..f3b96c50a53 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/tests/async-io-test.c++ @@ -0,0 +1,801 @@ +// Tests for kj-rs-io: tokio-backed implementations of KJ's async I/O interfaces, driven by a +// kj::EventLoop on a TokioEventPort. Following kj-rs conventions, C++ KJ_TESTs drive; Rust +// helpers (tests/lib.rs) provide the native-side behaviors (unwrap fast path, pre-bound fds). + +#include "kj-rs-io-test/lib.rs.h" +#include "kj-rs-io/async-io.h" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#if _WIN32 +#include // GetProcessTimes, for the wedge watchdog below. + +// After windows.h: un-breaks macros it leaks over KJ's, notably ERROR (which otherwise breaks +// the KJ_LOG(ERROR, ...) inside KJ_FAIL_* expansions). +#include +#else +#include +#include +#include +#include +#include +#include +#endif + +namespace kj_rs_io_test { +namespace { + +using kj_rs_io::setupTokioAsyncIo; +using kj_rs_io::TokioAsyncIoContext; + +// ======================================================================================= +// Helpers + +struct ConnectedPair { + kj::Own listener; + kj::Own client; + kj::Own server; +}; + +kj::Own parseNow( + TokioAsyncIoContext &io, kj::StringPtr addr, kj::uint portHint = 0) { + return io.getNetwork().parseAddress(addr, portHint).wait(io.getWaitScope()); +} + +ConnectedPair makeTcpPair(TokioAsyncIoContext &io) { + auto &ws = io.getWaitScope(); + auto listener = parseNow(io, "127.0.0.1")->listen(); + auto connectAddr = parseNow(io, kj::str("127.0.0.1:", listener->getPort())); + auto acceptPromise = listener->accept(); + auto client = connectAddr->connect().wait(ws); + auto server = acceptPromise.wait(ws); + return ConnectedPair{kj::mv(listener), kj::mv(client), kj::mv(server)}; +} + +kj::Array makePatternedData(size_t size, kj::byte seed) { + auto data = kj::heapArray(size); + for (size_t i = 0; i < size; i++) { + data[i] = static_cast((i * 31 + seed) & 0xff); + } + return data; +} + +::rust::Slice toRust(kj::ArrayPtr data) { + return ::rust::Slice(data.begin(), data.size()); +} + +::rust::Vec toRustVec(kj::ArrayPtr data) { + ::rust::Vec vec; + vec.reserve(data.size()); + for (auto b: data) { + vec.push_back(b); + } + return vec; +} + +// Writes `data` to `out` in chunks (exercising write-all + backpressure). +kj::Promise pumpOut(kj::AsyncIoStream &out, kj::ArrayPtr data) { + constexpr size_t CHUNK = 64 * 1024; + size_t offset = 0; + while (offset < data.size()) { + size_t n = kj::min(CHUNK, data.size() - offset); + co_await out.write(data.slice(offset, offset + n)); + offset += n; + } +} + +// Reads exactly `expected.size()` bytes from `in` and verifies they match `expected`. +kj::Promise drainAndCheck(kj::AsyncIoStream &in, kj::ArrayPtr expected) { + auto buffer = kj::heapArray(expected.size()); + size_t total = co_await in.tryRead(buffer.begin(), buffer.size(), buffer.size()); + KJ_ASSERT(total == expected.size(), total, expected.size()); + KJ_ASSERT(memcmp(buffer.begin(), expected.begin(), expected.size()) == 0); +} + +// CPU seconds consumed by this process so far. Used by the wedge watchdog to distinguish a +// spinning event loop from one that is parked and never woken. +double processCpuSeconds() { +#if _WIN32 + // clock() is WALL time on Windows (CRT quirk), so use GetProcessTimes. + FILETIME creationTime, exitTime, kernelTime, userTime; + GetProcessTimes(GetCurrentProcess(), &creationTime, &exitTime, &kernelTime, &userTime); + auto toSeconds = [](const FILETIME &ft) { + return static_cast( + (static_cast(ft.dwHighDateTime) << 32) | ft.dwLowDateTime) * + 1e-7; + }; + return toSeconds(kernelTime) + toSeconds(userTime); +#else + return static_cast(clock()) / CLOCKS_PER_SEC; +#endif +} + +// Waits for `connectPromise` (a connect to a certainly-closed port) to fail and returns the +// exception. Instrumented for a wedge observed (flakily) on Windows CI, where such a connect +// neither succeeded nor failed and a bare wait() silently ate the whole binary's bazel timeout: +// +// - A 30 s KJ timer bounds the wait, so a lost connect-readiness wake fails the test with a +// message instead. (A pending timer also makes the event port park with a timeout rather than +// indefinitely, so if the wedge is a lost wake on an indefinite park, the timer tick itself +// recovers it -- the run then passes, which is a data point too: the one CI run carrying this +// bound passed while both runs with a bare wait() timed out.) +// - A watchdog thread aborts after 60 s in case the loop stops servicing even timers, reporting +// process CPU use to distinguish a spinning loop (high) from one parked without wakeups (~0). +kj::Exception expectConnectFailure( + TokioAsyncIoContext &io, kj::Promise> connectPromise) { + std::atomic done{false}; + std::thread watchdog([&done]() { + double cpuBefore = processCpuSeconds(); + for (int i = 0; i < 600; i++) { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + if (done.load()) return; + } + double cpuUsed = processCpuSeconds() - cpuBefore; + fprintf(stderr, + "expectConnectFailure watchdog: the event loop serviced neither the connect nor the 30s " + "timer for 60s; the process used %.1f CPU-seconds meanwhile (high = loop spinning, ~0 = " + "parked and never woken). Aborting rather than eating the bazel timeout.\n", + cpuUsed); + fflush(stderr); + abort(); + }); + KJ_DEFER({ + done.store(true); + watchdog.join(); + }); + + auto timeout = + io.getTimer().afterDelay(30 * kj::SECONDS).then([]() -> kj::Own { + KJ_FAIL_ASSERT("connect() to a closed port neither succeeded nor failed within 30s; " + "the connect-failure readiness wake was likely lost"); + }); + return KJ_ASSERT_NONNULL(kj::runCatchingExceptions([&]() { + connectPromise.exclusiveJoin(kj::mv(timeout)).wait(io.getWaitScope()); + }), + "connect() to a closed port unexpectedly succeeded"); +} + +// ======================================================================================= +// Stream contract + +KJ_TEST("tryRead waits for minBytes, then returns what is available up to " + "maxBytes") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pair = makeTcpPair(io); + + kj::byte buffer[16]; + + // Exactly-min: 3 bytes written, min 3 -> resolves with 3. + pair.client->write("abc"_kjb).wait(ws); + KJ_EXPECT(pair.server->tryRead(buffer, 3, sizeof(buffer)).wait(ws) == 3); + KJ_EXPECT(kj::ArrayPtr(buffer, 3) == "abc"_kjb); + + // Blocks until minBytes: 2 available < min 5 -> pending; 3 more arrive -> resolves with 5. + pair.client->write("de"_kjb).wait(ws); + auto readPromise = pair.server->tryRead(buffer, 5, sizeof(buffer)); + KJ_EXPECT(!readPromise.poll(ws)); + pair.client->write("fgh"_kjb).wait(ws); + KJ_EXPECT(readPromise.wait(ws) == 5); + KJ_EXPECT(kj::ArrayPtr(buffer, 5) == "defgh"_kjb); +} + +KJ_TEST("EOF before minBytes returns a short count; half-close keeps the other " + "direction usable") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pair = makeTcpPair(io); + + pair.client->write("ab"_kjb).wait(ws); + pair.client->shutdownWrite(); + + // EOF-before-min: only 2 bytes then FIN -> tryRead(min 5) resolves with 2. + kj::byte buffer[16]; + KJ_EXPECT(pair.server->tryRead(buffer, 5, sizeof(buffer)).wait(ws) == 2); + KJ_EXPECT(kj::ArrayPtr(buffer, 2) == "ab"_kjb); + // Subsequent reads keep reporting EOF. + KJ_EXPECT(pair.server->tryRead(buffer, 1, sizeof(buffer)).wait(ws) == 0); + + // Half-close: server -> client direction still works after client's shutdownWrite. + pair.server->write("reply"_kjb).wait(ws); + KJ_EXPECT(pair.client->tryRead(buffer, 5, sizeof(buffer)).wait(ws) == 5); + KJ_EXPECT(kj::ArrayPtr(buffer, 5) == "reply"_kjb); +} + +KJ_TEST("multi-megabyte transfers in both directions with concurrent read+write " + "per stream") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pair = makeTcpPair(io); + + constexpr size_t SIZE = 8 * 1024 * 1024; + auto dataA = makePatternedData(SIZE, 1); + auto dataB = makePatternedData(SIZE, 2); + + // All four directions at once: each stream is simultaneously reading and writing, and each + // transfer is far larger than the socket buffers (forcing many readiness round-trips). + auto builder = kj::heapArrayBuilder>(4); + builder.add(pumpOut(*pair.client, dataA)); + builder.add(drainAndCheck(*pair.server, dataA)); + builder.add(pumpOut(*pair.server, dataB)); + builder.add(drainAndCheck(*pair.client, dataB)); + kj::joinPromisesFailFast(builder.finish()).wait(ws); +} + +KJ_TEST("multi-piece write() writes all pieces in order") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pair = makeTcpPair(io); + + const kj::ArrayPtr pieces[] = {"one,"_kjb, "two,"_kjb, "three"_kjb}; + pair.client->write(kj::arrayPtr(pieces, 3)).wait(ws); + + kj::byte buffer[32]; + KJ_EXPECT(pair.server->tryRead(buffer, 13, sizeof(buffer)).wait(ws) == 13); + KJ_EXPECT(kj::ArrayPtr(buffer, 13) == "one,two,three"_kjb); +} + +KJ_TEST("canceling a blocked read releases the socket for reuse") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pair = makeTcpPair(io); + + kj::byte buffer[16]; + { + // A read blocked in tokio (registered with the I/O driver, no data available)... + auto blocked = pair.server->tryRead(buffer, 1, sizeof(buffer)); + KJ_EXPECT(!blocked.poll(ws)); + // ...is canceled by dropping the promise, which must drop the Rust future and release the + // read interest. + } + + // The stream remains fully usable: a fresh read gets the next bytes. + pair.client->write("later"_kjb).wait(ws); + KJ_EXPECT(pair.server->tryRead(buffer, 5, sizeof(buffer)).wait(ws) == 5); + KJ_EXPECT(kj::ArrayPtr(buffer, 5) == "later"_kjb); + + // Canceling mid-large-write also leaves the process sane (bytes may be lost, like KJ). + { + auto data = makePatternedData(16 * 1024 * 1024, 7); + auto bigWrite = pair.client->write(data); + if (bigWrite.poll(ws)) { + bigWrite.wait(ws); + } + } +} + +#if !_WIN32 +KJ_TEST("whenWriteDisconnected resolves on peer reset, not on half-close") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pair = makeTcpPair(io); + + auto disconnected = pair.client->whenWriteDisconnected(); + KJ_EXPECT(!disconnected.poll(ws)); + + // A peer half-close (FIN) must NOT count as write-disconnect: the client can still write. + pair.server->shutdownWrite(); + kj::byte buffer[16]; + KJ_EXPECT(pair.client->tryRead(buffer, 1, sizeof(buffer)).wait(ws) == 0); // observe EOF + KJ_EXPECT(!disconnected.poll(ws)); + + // Destroying the server end with SO_LINGER=0 sends an RST; now writes are doomed. + struct linger lin; + lin.l_onoff = 1; + lin.l_linger = 0; + pair.server->setsockopt(SOL_SOCKET, SO_LINGER, &lin, sizeof(lin)); + pair.server = nullptr; + + disconnected.wait(ws); +} +#endif + +KJ_TEST("acceptAuthenticated reports the TCP peer's NetworkPeerIdentity") { + // workerd's HTTP listener builds the cf blob's clientIp (-> the CF-Connecting-IP header) from + // this identity; UnknownPeerIdentity (kj's base-class default) silently yields an empty + // client IP. + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + auto listener = parseNow(io, "127.0.0.1")->listen(); + auto acceptPromise = listener->acceptAuthenticated(); + auto client = parseNow(io, kj::str("127.0.0.1:", listener->getPort()))->connect().wait(ws); + auto server = acceptPromise.wait(ws); + + auto &identity = + KJ_ASSERT_NONNULL(kj::tryDowncast(*server.peerIdentity)); + // KJ's "ip:port" format, byte-identical to the native backend. + auto text = identity.toString(); + KJ_EXPECT(text.startsWith("127.0.0.1:"), text); +} + +#if !_WIN32 +KJ_TEST("acceptAuthenticated reports LocalPeerIdentity credentials on unix sockets") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + // /tmp rather than TEST_TMPDIR: sun_path is limited to ~104 bytes. + auto path = kj::str("/tmp/kj-rs-io-auth-test-", getpid(), ".sock"); + auto addr = parseNow(io, kj::str("unix:", path)); + + auto listener = addr->listen(); + auto acceptPromise = listener->acceptAuthenticated(); + auto client = addr->connect().wait(ws); + auto server = acceptPromise.wait(ws); + + auto &identity = KJ_ASSERT_NONNULL(kj::tryDowncast(*server.peerIdentity)); + auto creds = identity.getCredentials(); + // The peer is this very process. + KJ_EXPECT(KJ_ASSERT_NONNULL(creds.pid) == getpid()); + KJ_EXPECT(KJ_ASSERT_NONNULL(creds.uid) == getuid()); + + unlink(path.cStr()); +} +#endif + +KJ_TEST("sockname/peername/sockopt/getFd passthrough") { +#if _WIN32 + return; +#else + auto io = setupTokioAsyncIo(); + auto pair = makeTcpPair(io); + + // getFd is populated. + KJ_EXPECT(KJ_ASSERT_NONNULL(pair.client->getFd()) >= 0); + + // The client's peer is the server's local socket. + struct sockaddr_in peer, local; + kj::uint peerLen = sizeof(peer), localLen = sizeof(local); + pair.client->getpeername(reinterpret_cast(&peer), &peerLen); + pair.server->getsockname(reinterpret_cast(&local), &localLen); + KJ_EXPECT(peer.sin_port == local.sin_port); + KJ_EXPECT(peer.sin_addr.s_addr == local.sin_addr.s_addr); + + // setsockopt/getsockopt round-trip (this is also how setNoDelay-style options are applied). + int on = 1; + pair.client->setsockopt(IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on)); + int result = 0; + kj::uint resultLen = sizeof(result); + pair.client->getsockopt(IPPROTO_TCP, TCP_NODELAY, &result, &resultLen); + KJ_EXPECT(result != 0); +#endif +} + +// ======================================================================================= +// Network / addresses + +KJ_TEST("parseAddress handles IP literals, port hints, and toString round-trips") { + auto io = setupTokioAsyncIo(); + + KJ_EXPECT(parseNow(io, "1.2.3.4:80")->toString() == "1.2.3.4:80"); + KJ_EXPECT(parseNow(io, "1.2.3.4", 99)->toString() == "1.2.3.4:99"); + KJ_EXPECT(parseNow(io, "[1234:5678::abcd]:80")->toString() == "[1234:5678::abcd]:80"); + KJ_EXPECT(parseNow(io, "1234:5678::abcd", 80)->toString() == "[1234:5678::abcd]:80"); + KJ_EXPECT(parseNow(io, "*:80")->toString() == "*:80"); + KJ_EXPECT(parseNow(io, "*")->toString() == "*:0"); + + // clone() produces an equivalent address. + auto addr = parseNow(io, "127.0.0.1:1234"); + KJ_EXPECT(addr->clone()->toString() == addr->toString()); +} + +KJ_TEST("wildcard listen binds dual-stack and reports its port; port 0 picks a " + "free port") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + auto listener = parseNow(io, "*:0")->listen(); + kj::uint port = listener->getPort(); + KJ_EXPECT(port != 0); + + // Reachable over both IPv4 and IPv6 loopback (IPV6_V6ONLY off, like KJ). + kj::String addrTexts[] = {kj::str("127.0.0.1:", port), kj::str("[::1]:", port)}; + for (auto &addrText: addrTexts) { + auto acceptPromise = listener->accept(); + auto client = parseNow(io, addrText)->connect().wait(ws); + auto server = acceptPromise.wait(ws); + client->write("ping"_kjb).wait(ws); + kj::byte buffer[4]; + KJ_EXPECT(server->tryRead(buffer, 4, sizeof(buffer)).wait(ws) == 4); + } +} + +KJ_TEST("parseAddress resolves hostnames via DNS and connect tries addresses in " + "order") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + // Listen on IPv4 loopback only. "localhost" typically resolves to both ::1 and 127.0.0.1; + // connect() must try each in order until one succeeds. + auto listener = parseNow(io, "127.0.0.1")->listen(); + auto addr = parseNow(io, kj::str("localhost:", listener->getPort())); + + auto acceptPromise = listener->accept(); + auto client = addr->connect().wait(ws); + auto server = acceptPromise.wait(ws); + client->write("dns!"_kjb).wait(ws); + kj::byte buffer[4]; + KJ_EXPECT(server->tryRead(buffer, 4, sizeof(buffer)).wait(ws) == 4); + KJ_EXPECT(kj::ArrayPtr(buffer, 4) == "dns!"_kjb); +} + +KJ_TEST("connect to a closed port surfaces a DISCONNECTED kj::Exception " + "mentioning the refusal") { + auto io = setupTokioAsyncIo(); + + // Find a port that is certainly closed: bind one, note it, close it. + kj::uint port; + { + auto listener = parseNow(io, "127.0.0.1")->listen(); + port = listener->getPort(); + } + + auto addr = parseNow(io, kj::str("127.0.0.1:", port)); + auto exception = expectConnectFailure(io, addr->connect()); + // Exact text (recorded): "connect(): Connection refused (os error 61)" on macOS / + // "... (os error 111)" on Linux. KJ's native text would be "connect(): Connection refused". + KJ_EXPECT( + strstr(exception.getDescription().cStr(), "refused") != nullptr, exception.getDescription()); + KJ_EXPECT(exception.getType() == kj::Exception::Type::DISCONNECTED); +} + +KJ_TEST("the address may be dropped while connect() is pending (KJ lifetime " + "contract)") { + // Upstream KJ heap-copies the resolved address list into the connect promise + // (NetworkAddressImpl::connect() in kj/async-io-unix.c++), so callers may legally drop + // the kj::NetworkAddress right after calling connect(). Verify this port honors the same + // contract, on both the success path and the error/retry path. + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + // Success path: drop the address immediately, then complete the connect. + { + auto listener = parseNow(io, "127.0.0.1")->listen(); + auto acceptPromise = listener->accept(); + + kj::Promise> connectPromise = nullptr; + { + auto addr = parseNow(io, kj::str("127.0.0.1:", listener->getPort())); + connectPromise = addr->connect(); + // `addr` is destroyed here, while the connect is still in flight. + } + + auto client = connectPromise.wait(ws); + auto server = acceptPromise.wait(ws); + client->write("hello"_kjb).wait(ws); + kj::byte buffer[5] = {}; + KJ_EXPECT(server->tryRead(buffer, 5, 5).wait(ws) == 5); + KJ_EXPECT(kj::arrayPtr(buffer, 5) == "hello"_kjb); + } + + // Error path: connect to a certainly-closed port with the address already dropped; the + // failure continuation (which re-reads the address list) must still be safe and surface + // the normal exception. + { + kj::uint port; + { + auto listener = parseNow(io, "127.0.0.1")->listen(); + port = listener->getPort(); + } + + kj::Promise> connectPromise = nullptr; + { + auto addr = parseNow(io, kj::str("127.0.0.1:", port)); + connectPromise = addr->connect(); + } + + auto exception = expectConnectFailure(io, kj::mv(connectPromise)); + KJ_EXPECT(strstr(exception.getDescription().cStr(), "refused") != nullptr, + exception.getDescription()); + KJ_EXPECT(exception.getType() == kj::Exception::Type::DISCONNECTED); + } +} + +KJ_TEST("connecting to a wildcard address is an error") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto addr = parseNow(io, "*:1234"); + KJ_EXPECT_THROW_MESSAGE("wildcard", addr->connect().wait(ws)); +} + +#if !_WIN32 +KJ_TEST("unix domain sockets: parse, listen, connect, toString") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + // Note: /tmp rather than TEST_TMPDIR because sun_path is limited to ~104 bytes. + auto path = kj::str("/tmp/kj-rs-io-test-", getpid(), ".sock"); + auto addrText = kj::str("unix:", path); + + auto addr = parseNow(io, addrText); + KJ_EXPECT(addr->toString() == addrText); + + auto listener = addr->listen(); + KJ_EXPECT(listener->getPort() == 0); // KJ reports 0 for non-IP listeners. + auto acceptPromise = listener->accept(); + auto client = addr->connect().wait(ws); + auto server = acceptPromise.wait(ws); + + client->write("via unix"_kjb).wait(ws); + client->shutdownWrite(); + kj::byte buffer[16]; + KJ_EXPECT(server->tryRead(buffer, 16, sizeof(buffer)).wait(ws) == 8); + KJ_EXPECT(kj::ArrayPtr(buffer, 8) == "via unix"_kjb); + + unlink(path.cStr()); +} + +KJ_TEST("getSockaddr builds a connectable address from a raw struct sockaddr") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + auto listener = parseNow(io, "127.0.0.1")->listen(); + + struct sockaddr_in sin; + memset(&sin, 0, sizeof(sin)); + sin.sin_family = AF_INET; + sin.sin_port = htons(static_cast(listener->getPort())); + sin.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + auto addr = io.getNetwork().getSockaddr(&sin, sizeof(sin)); + KJ_EXPECT(addr->toString() == kj::str("127.0.0.1:", listener->getPort())); + + auto acceptPromise = listener->accept(); + auto client = addr->connect().wait(ws); + auto server = acceptPromise.wait(ws); + client->write("hi"_kjb).wait(ws); + kj::byte buffer[2]; + KJ_EXPECT(server->tryRead(buffer, 2, sizeof(buffer)).wait(ws) == 2); +} +#endif + +KJ_TEST("newPipeThread is a documented stub") { + auto io = setupTokioAsyncIo(); + KJ_EXPECT_THROW_MESSAGE("newPipeThread", + io.getProvider().newPipeThread( + [](kj::AsyncIoProvider &, kj::AsyncIoStream &, kj::WaitScope &) {})); +} + +// ======================================================================================= +// Provider odds and ends + +KJ_TEST("provider pipes (in-memory) and timer work under the tokio loop") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + auto pipe = io.getProvider().newTwoWayPipe(); + auto writePromise = pipe.ends[0]->write("pipe data"_kjb).eagerlyEvaluate(nullptr); + kj::byte buffer[16]; + KJ_EXPECT(pipe.ends[1]->tryRead(buffer, 9, sizeof(buffer)).wait(ws) == 9); + writePromise.wait(ws); + + auto &timer = io.getProvider().getTimer(); + auto before = timer.now(); + timer.afterDelay(5 * kj::MILLISECONDS).wait(ws); + KJ_EXPECT(timer.now() - before >= 5 * kj::MILLISECONDS); +} + +// ======================================================================================= +// Unwrap fast path + +KJ_TEST("unwrap fast path: recover the native tokio stream and write from Rust") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pair = makeTcpPair(io); + + // Recover the native tokio TcpStream out of the kj wrapper (free-function form, as a Rust + // server would after being handed a kj::AsyncIoStream&)... + auto native = kj_rs_io::unwrapTokioStream(*pair.server); + + // ...the hollow wrapper now refuses I/O... + kj::byte buffer[32]; + KJ_EXPECT_THROW_MESSAGE("unwrapped", pair.server->tryRead(buffer, 1, sizeof(buffer)).wait(ws)); + + // ...and Rust can drive the connection natively: it writes via the tokio readiness API and + // closes; C++ reads the bytes plus EOF through the (still wrapped) client end. + auto writeDone = native_write_via_unwrap(kj::mv(native), toRustVec("native write"_kjb)); + KJ_EXPECT(pair.client->tryRead(buffer, 12, sizeof(buffer)).wait(ws) == 12); + KJ_EXPECT(kj::ArrayPtr(buffer, 12) == "native write"_kjb); + writeDone.wait(ws); + KJ_EXPECT(pair.client->tryRead(buffer, 1, sizeof(buffer)).wait(ws) == 0); // EOF +} + +KJ_TEST("unwrap fast path: Rust-side unwrap_kj_stream() from a " + "kj::AsyncIoStream&") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pair = makeTcpPair(io); + + // Rust receives only a kj::AsyncIoStream& and performs the unwrap + native write itself. + auto writeDone = native_write_via_kj_unwrap(*pair.server, toRust("rust unwrap"_kjb)); + kj::byte buffer[32]; + KJ_EXPECT(pair.client->tryRead(buffer, 11, sizeof(buffer)).wait(ws) == 11); + KJ_EXPECT(kj::ArrayPtr(buffer, 11) == "rust unwrap"_kjb); + writeDone.wait(ws); + + // Unwrapping a foreign (non-kj-rs-io) stream fails cleanly. + auto pipe = io.getProvider().newTwoWayPipe(); + KJ_EXPECT_THROW_MESSAGE("cannot unwrap", kj_rs_io::unwrapTokioStream(*pipe.ends[0])); +} + +// ======================================================================================= +// Fd wrapping (kj::LowLevelAsyncIoProvider) + +#if !_WIN32 +KJ_TEST("wrapListenSocketFd accepts connections on a pre-bound listener (the " + "--socket-fd case)") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + // Rust binds a *blocking* std listener (like an fd inherited from a supervisor) and hands us + // the raw fd; wrapListenSocketFd must take ownership and make it usable. + auto prebound = create_prebound_listener_fd(); + auto receiver = io.getLowLevelProvider().wrapListenSocketFd( + prebound.fd, kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP); + KJ_EXPECT(receiver->getPort() == prebound.port); + + auto acceptPromise = receiver->accept(); + auto client = parseNow(io, kj::str("127.0.0.1:", prebound.port))->connect().wait(ws); + auto server = acceptPromise.wait(ws); + + client->write("fd listen"_kjb).wait(ws); + kj::byte buffer[16]; + KJ_EXPECT(server->tryRead(buffer, 9, sizeof(buffer)).wait(ws) == 9); + KJ_EXPECT(kj::ArrayPtr(buffer, 9) == "fd listen"_kjb); +} + +KJ_TEST("wrapSocketFd wraps both ends of a socketpair") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + int fds[2]; + KJ_SYSCALL(socketpair(AF_UNIX, SOCK_STREAM, 0, fds)); + auto end0 = + io.getLowLevelProvider().wrapSocketFd(fds[0], kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP); + auto end1 = + io.getLowLevelProvider().wrapSocketFd(fds[1], kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP); + + end0->write("socketpair"_kjb).wait(ws); + kj::byte buffer[16]; + KJ_EXPECT(end1->tryRead(buffer, 10, sizeof(buffer)).wait(ws) == 10); + KJ_EXPECT(kj::ArrayPtr(buffer, 10) == "socketpair"_kjb); +} + +KJ_TEST("wrapConnectingSocketFd completes a nonblocking connect") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + auto listener = parseNow(io, "127.0.0.1")->listen(); + + int fd; + KJ_SYSCALL(fd = socket(AF_INET, SOCK_STREAM, 0)); + struct sockaddr_in sin; + memset(&sin, 0, sizeof(sin)); + sin.sin_family = AF_INET; + sin.sin_port = htons(static_cast(listener->getPort())); + sin.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + + auto acceptPromise = listener->accept(); + auto client = io.getLowLevelProvider() + .wrapConnectingSocketFd(fd, reinterpret_cast(&sin), + sizeof(sin), kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP) + .wait(ws); + auto server = acceptPromise.wait(ws); + + client->write("connected"_kjb).wait(ws); + kj::byte buffer[16]; + KJ_EXPECT(server->tryRead(buffer, 9, sizeof(buffer)).wait(ws) == 9); +} + +KJ_TEST("wrapInputFd/wrapOutputFd move bytes through an OS pipe and observe EOF") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + int fds[2]; + KJ_SYSCALL(pipe(fds)); + auto input = + io.getLowLevelProvider().wrapInputFd(fds[0], kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP); + auto output = + io.getLowLevelProvider().wrapOutputFd(fds[1], kj::LowLevelAsyncIoProvider::TAKE_OWNERSHIP); + + // Blocked read completes once data is written ("pipe fd io" is 10 bytes; cap the first read + // at 9 so one byte remains). + kj::byte buffer[16]; + auto readPromise = input->tryRead(buffer, 9, 9); + KJ_EXPECT(!readPromise.poll(ws)); + auto writePromise = output->write("pipe fd io"_kjb); + KJ_EXPECT(readPromise.wait(ws) == 9); + writePromise.wait(ws); + KJ_EXPECT(input->tryRead(buffer, 1, sizeof(buffer)).wait(ws) == 1); // "o" + + // Dropping the output stream closes the write end -> EOF. + output = nullptr; + KJ_EXPECT(input->tryRead(buffer, 1, sizeof(buffer)).wait(ws) == 0); +} + +KJ_TEST("restrictPeers blocks disallowed connect() with KJ's error text") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + kj_rs_io::TokioNetwork network; + auto restricted = network.restrictPeers({"public"_kj}, {}); + + // Loopback is not "public": blocked before any connection attempt. + auto blockedAddr = restricted->parseAddress("127.0.0.1:1").wait(ws); + KJ_EXPECT_THROW_MESSAGE("connect() blocked by restrictPeers()", blockedAddr->connect().wait(ws)); + + // getSockaddr is rejected eagerly, like KJ. + struct sockaddr_in sin; + memset(&sin, 0, sizeof(sin)); + sin.sin_family = AF_INET; + sin.sin_port = htons(1); + sin.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + KJ_EXPECT_THROW_MESSAGE( + "address blocked by restrictPeers()", restricted->getSockaddr(&sin, sizeof(sin))); + + // An allowing restriction still connects. + auto allowed = network.restrictPeers({"private"_kj}, {}); + auto listener = network.parseAddress("127.0.0.1", 0).wait(ws)->listen(); + auto acceptPromise = listener->accept(); + auto client = allowed->parseAddress(kj::str("127.0.0.1:", listener->getPort())) + .wait(ws) + ->connect() + .wait(ws); + auto server = acceptPromise.wait(ws); + client->write("ok"_kjb).wait(ws); + kj::byte buffer[2]; + KJ_EXPECT(server->tryRead(buffer, 2, 2).wait(ws) == 2); +} + +KJ_TEST("restrictPeers filters accepted peers (disallowed peers are dropped, " + "accept keeps waiting)") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + kj_rs_io::TokioNetwork network; + auto restricted = network.restrictPeers({"public"_kj}, {}); + + auto listener = restricted->parseAddress("127.0.0.1", 0).wait(ws)->listen(); + auto acceptPromise = listener->accept(); + + // Connect via the unrestricted network; the loopback peer is not "public", so the listener + // silently drops it: accept() stays pending and the client observes EOF. + auto client = network.parseAddress(kj::str("127.0.0.1:", listener->getPort()), 0) + .wait(ws) + ->connect() + .wait(ws); + KJ_EXPECT(!acceptPromise.poll(ws)); + kj::byte buffer[1]; + KJ_EXPECT(client->tryRead(buffer, 1, 1).wait(ws) == 0); +} + +KJ_TEST("onSignal resolves when the process receives the signal") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + auto promise = kj_rs_io::onSignal(SIGUSR2); + // Pump the loop so the (cold, first-poll-registered) tokio signal handler is installed + // before we raise; raising first would take SIGUSR2's default disposition (terminate). + KJ_EXPECT(!promise.poll(ws)); + + KJ_SYSCALL(kill(getpid(), SIGUSR2)); + promise.wait(ws); + + // A second watcher works too (the process-global registration is reusable). + auto again = kj_rs_io::onSignal(SIGUSR2); + KJ_EXPECT(!again.poll(ws)); + KJ_SYSCALL(kill(getpid(), SIGUSR2)); + again.wait(ws); +} +#endif + +} // namespace +} // namespace kj_rs_io_test diff --git a/src/rust/cxx/kj-rs-io/tests/capnp-rpc-test.c++ b/src/rust/cxx/kj-rs-io/tests/capnp-rpc-test.c++ new file mode 100644 index 00000000000..b1ea271c8b3 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/tests/capnp-rpc-test.c++ @@ -0,0 +1,79 @@ +// capnp RPC integration: capnp::TwoPartyServer / capnp::TwoPartyClient running over kj-rs-io +// (tokio-backed) streams on the tokio event loop — bootstrap plus call round-trips. Uses a +// schema-less capability (raw dispatchCall / typelessRequest with Text payloads) to avoid +// needing capnp codegen in this repo. + +#include "kj-rs-io/async-io.h" + +#include +#include +#include +#include +#include +#include + +namespace kj_rs_io_test { +namespace { + +using kj_rs_io::setupTokioAsyncIo; + +constexpr uint64_t ECHO_INTERFACE_ID = 0xabcd1234abcd1234ull; +constexpr uint16_t ECHO_METHOD_ID = 0; + +// A schema-less capability: method 0 takes a Text param and returns "echo:" + text. +class EchoCapability final: public capnp::Capability::Server { + public: + DispatchCallResult dispatchCall(uint64_t interfaceId, + uint16_t methodId, + capnp::CallContext context) override { + KJ_ASSERT(interfaceId == ECHO_INTERFACE_ID); + KJ_ASSERT(methodId == ECHO_METHOD_ID); + auto params = kj::str(context.getParams().getAs()); + context.releaseParams(); + context.getResults(capnp::MessageSize{16, 0}).setAs(kj::str("echo:", params)); + return DispatchCallResult{kj::READY_NOW, false, true}; + } +}; + +KJ_TEST("capnp two-party RPC bootstrap and call round-trips over kj-rs-io " + "streams on the tokio loop") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + // Server: TwoPartyServer accepting from a kj-rs-io ConnectionReceiver. + capnp::TwoPartyServer server(kj::heap()); + auto listener = io.getNetwork().parseAddress("127.0.0.1", 0).wait(ws)->listen(); + auto listenTask = server.listen(*listener).eagerlyEvaluate( + [](kj::Exception &&e) { KJ_FAIL_EXPECT("RPC server failed", e); }); + + // Client: TwoPartyClient over a kj-rs-io connection. + auto addr = io.getNetwork().parseAddress(kj::str("127.0.0.1:", listener->getPort())).wait(ws); + auto connection = addr->connect().wait(ws); + capnp::TwoPartyClient client(*connection); + auto cap = client.bootstrap(); + + // Single call round trip. + { + auto request = cap.typelessRequest(ECHO_INTERFACE_ID, ECHO_METHOD_ID, kj::none, {}); + request.setAs("hello tokio"); + auto response = request.send().wait(ws); + KJ_EXPECT(response.getAs() == "echo:hello tokio"); + } + + // A pile of pipelined calls in flight at once. + { + constexpr int COUNT = 64; + auto builder = kj::heapArrayBuilder>(COUNT); + for (int i = 0; i < COUNT; i++) { + auto request = cap.typelessRequest(ECHO_INTERFACE_ID, ECHO_METHOD_ID, kj::none, {}); + request.setAs(kj::str("msg", i)); + builder.add(request.send().then([i](capnp::Response response) { + KJ_EXPECT(response.getAs() == kj::str("echo:msg", i)); + })); + } + kj::joinPromisesFailFast(builder.finish()).wait(ws); + } +} + +} // namespace +} // namespace kj_rs_io_test diff --git a/src/rust/cxx/kj-rs-io/tests/file-watcher-test.c++ b/src/rust/cxx/kj-rs-io/tests/file-watcher-test.c++ new file mode 100644 index 00000000000..b3ad24168e6 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/tests/file-watcher-test.c++ @@ -0,0 +1,274 @@ +// Tests for kj_rs_io::FileWatcher (file-watcher.h), the tokio-loop replacement for workerd's +// --watch FileWatcher. Exercises the behaviors workerd depends on: plain modification, atomic +// replace-by-rename (editor saves), event queueing/coalescing across onChange() calls, the +// already-open-fd watch path (kqueue backends), missing-file handling, and teardown/cancel +// while a watch promise is armed. + +#include "kj-rs-io/async-io.h" +#include "kj-rs-io/file-watcher.h" + +#include +#include +#include +#include + +#include +#include + +#if !_WIN32 +#include +#include +#include +#endif + +namespace kj_rs_io_test { +namespace { + +using kj_rs_io::FileWatcher; +using kj_rs_io::setupTokioAsyncIo; +using kj_rs_io::TokioAsyncIoContext; + +#if !_WIN32 + +// ======================================================================================= +// Helpers + +// Waits for `promise` to resolve, returning true, or false after `timeout`. +bool resolvesWithin(kj::Promise promise, TokioAsyncIoContext &io, kj::Duration timeout) { + auto timedOut = io.getTimer().afterDelay(timeout).then([]() { return false; }); + return promise.then([]() { return true; }) + .exclusiveJoin(kj::mv(timedOut)) + .wait(io.getWaitScope()); +} + +// Generous bound for "the change fires"; file events are near-immediate on both backends. +constexpr kj::Duration FIRE_TIMEOUT = 5 * kj::SECONDS; +// Short bound for "nothing fires" checks. +constexpr kj::Duration QUIET_TIMEOUT = 200 * kj::MILLISECONDS; + +struct TempDir { + kj::String path; + + TempDir() { + const char *base = getenv("TEST_TMPDIR"); + if (base == nullptr) base = "/tmp"; + auto tmpl = kj::str(base, "/kj-rs-io-file-watcher-test.XXXXXX"); + KJ_ASSERT(mkdtemp(tmpl.begin()) != nullptr, strerror(errno)); + path = kj::mv(tmpl); + } + + ~TempDir() noexcept(false) { + // Best-effort cleanup; TEST_TMPDIR is wiped by bazel anyway. + auto cmd = kj::str("rm -rf ", path); + (void)system(cmd.cStr()); + } + + kj::String fileName(kj::StringPtr name) { + return kj::str(path, "/", name); + } + + kj::Path filePath(kj::StringPtr name) { + auto full = fileName(name); + KJ_ASSERT(full.startsWith("/")); + return kj::Path::parse(full.slice(1)); + } +}; + +void writeFile(kj::StringPtr path, kj::StringPtr content) { + kj::OwnFd fd = KJ_SYSCALL_FD(open(path.cStr(), O_WRONLY | O_CREAT | O_TRUNC, 0644)); + KJ_SYSCALL(write(fd, content.begin(), content.size())); +} + +void appendFile(kj::StringPtr path, kj::StringPtr content) { + kj::OwnFd fd = KJ_SYSCALL_FD(open(path.cStr(), O_WRONLY | O_APPEND)); + KJ_SYSCALL(write(fd, content.begin(), content.size())); +} + +// After a change fired, drains any further already-queued events so the next onChange() call +// starts from a quiet state (mirrors what workerd's waitForChanges() settle loop achieves). +void drain(FileWatcher &watcher, TokioAsyncIoContext &io) { + while (resolvesWithin(watcher.onChange(), io, QUIET_TIMEOUT)) {} +} + +// ======================================================================================= +// Tests + +KJ_TEST("FileWatcher: supported on this platform") { + auto io = setupTokioAsyncIo(); + FileWatcher watcher; + KJ_EXPECT(watcher.isSupported()); +} + +KJ_TEST("FileWatcher: modification fires onChange") { + auto io = setupTokioAsyncIo(); + TempDir dir; + writeFile(dir.fileName("a.txt"), "one"); + + FileWatcher watcher; + watcher.watch(dir.filePath("a.txt"), kj::none); + + auto change = watcher.onChange(); + appendFile(dir.fileName("a.txt"), " two"); + KJ_EXPECT(resolvesWithin(kj::mv(change), io, FIRE_TIMEOUT)); +} + +KJ_TEST("FileWatcher: change before onChange() is called is not lost") { + auto io = setupTokioAsyncIo(); + TempDir dir; + writeFile(dir.fileName("a.txt"), "one"); + + FileWatcher watcher; + watcher.watch(dir.filePath("a.txt"), kj::none); + + // Modify before anyone is waiting: the event queues in the kernel. + appendFile(dir.fileName("a.txt"), " two"); + KJ_EXPECT(resolvesWithin(watcher.onChange(), io, FIRE_TIMEOUT)); +} + +KJ_TEST("FileWatcher: atomic replace-by-rename fires onChange") { + // Editors typically save by writing a temporary file and rename(2)ing it over the target. + auto io = setupTokioAsyncIo(); + TempDir dir; + writeFile(dir.fileName("a.txt"), "one"); + + FileWatcher watcher; + watcher.watch(dir.filePath("a.txt"), kj::none); + + auto change = watcher.onChange(); + writeFile(dir.fileName("a.txt.tmp"), "two"); + KJ_SYSCALL(rename(dir.fileName("a.txt.tmp").cStr(), dir.fileName("a.txt").cStr())); + KJ_EXPECT(resolvesWithin(kj::mv(change), io, FIRE_TIMEOUT)); +} + +KJ_TEST("FileWatcher: watching via an already-open file handle") { + // workerd passes the config files' already-open kj::ReadableFile to watch(); the kqueue + // backend watches a dup of that fd (the inotify backend ignores it and uses the path). + auto io = setupTokioAsyncIo(); + TempDir dir; + writeFile(dir.fileName("a.txt"), "one"); + + auto file = kj::newDiskReadableFile(KJ_SYSCALL_FD(open(dir.fileName("a.txt").cStr(), O_RDONLY))); + + FileWatcher watcher; + watcher.watch(dir.filePath("a.txt"), *file); + file = nullptr; // The original handle may be closed; the watch must survive. + + auto change = watcher.onChange(); + appendFile(dir.fileName("a.txt"), " two"); + KJ_EXPECT(resolvesWithin(kj::mv(change), io, FIRE_TIMEOUT)); +} + +KJ_TEST("FileWatcher: rapid changes coalesce; watcher stays armed for later " + "changes") { + auto io = setupTokioAsyncIo(); + TempDir dir; + writeFile(dir.fileName("a.txt"), "one"); + + FileWatcher watcher; + watcher.watch(dir.filePath("a.txt"), kj::none); + + // A burst of changes produces one resolution per onChange() call (not one per event), ... + auto change = watcher.onChange(); + appendFile(dir.fileName("a.txt"), " two"); + appendFile(dir.fileName("a.txt"), " three"); + appendFile(dir.fileName("a.txt"), " four"); + KJ_EXPECT(resolvesWithin(kj::mv(change), io, FIRE_TIMEOUT)); + + // ... and once the queue is drained, the watcher is quiet ... + drain(watcher, io); + + // ... but still armed: a fresh change fires a fresh onChange(). + auto later = watcher.onChange(); + appendFile(dir.fileName("a.txt"), " five"); + KJ_EXPECT(resolvesWithin(kj::mv(later), io, FIRE_TIMEOUT)); +} + +KJ_TEST("FileWatcher: unrelated files in the same directory do not fire " + "(inotify filtering)") { + auto io = setupTokioAsyncIo(); + TempDir dir; + writeFile(dir.fileName("a.txt"), "one"); + writeFile(dir.fileName("other.txt"), "other"); + + FileWatcher watcher; + watcher.watch(dir.filePath("a.txt"), kj::none); + + appendFile(dir.fileName("other.txt"), " more"); + KJ_EXPECT(!resolvesWithin(watcher.onChange(), io, QUIET_TIMEOUT)); +} + +#if __linux__ +KJ_TEST("FileWatcher: watching a not-yet-existing file fires when it is created") { + // The inotify backend watches the parent directory, so the file itself need not exist yet. + // (The kqueue backend opens the file and so requires it to exist; see the test below.) + auto io = setupTokioAsyncIo(); + TempDir dir; + + FileWatcher watcher; + watcher.watch(dir.filePath("missing.txt"), kj::none); + + auto change = watcher.onChange(); + writeFile(dir.fileName("missing.txt"), "now it exists"); + KJ_EXPECT(resolvesWithin(kj::mv(change), io, FIRE_TIMEOUT)); +} +#else +KJ_TEST("FileWatcher: watching a nonexistent file throws (kqueue backend)") { + // Same behavior as workerd's kj-mode kqueue watcher: watch() opens the path with + // KJ_SYSCALL, which throws if it doesn't exist. + auto io = setupTokioAsyncIo(); + TempDir dir; + + FileWatcher watcher; + auto exception = + kj::runCatchingExceptions([&]() { watcher.watch(dir.filePath("missing.txt"), kj::none); }); + KJ_EXPECT(exception != kj::none); +} +#endif + +KJ_TEST("FileWatcher: canceling an armed onChange() and re-arming works") { + auto io = setupTokioAsyncIo(); + TempDir dir; + writeFile(dir.fileName("a.txt"), "one"); + + FileWatcher watcher; + watcher.watch(dir.filePath("a.txt"), kj::none); + + { + auto armed = watcher.onChange(); + KJ_EXPECT(!armed.poll(io.getWaitScope())); + // Dropped here while armed (fd registered with the tokio I/O driver). + } + + auto change = watcher.onChange(); + appendFile(dir.fileName("a.txt"), " two"); + KJ_EXPECT(resolvesWithin(kj::mv(change), io, FIRE_TIMEOUT)); +} + +KJ_TEST("FileWatcher: teardown while a watch promise is armed") { + auto io = setupTokioAsyncIo(); + TempDir dir; + writeFile(dir.fileName("a.txt"), "one"); + + auto watcher = kj::heap(); + watcher->watch(dir.filePath("a.txt"), kj::none); + + auto armed = watcher->onChange(); + KJ_EXPECT(!armed.poll(io.getWaitScope())); + + // Promise first (it borrows the watcher's fd), then the watcher itself. + armed = nullptr; + watcher = nullptr; +} + +#else // _WIN32 + +KJ_TEST("FileWatcher: reports unsupported on this platform") { + auto io = setupTokioAsyncIo(); + FileWatcher watcher; + KJ_EXPECT(!watcher.isSupported()); +} + +#endif + +} // namespace +} // namespace kj_rs_io_test diff --git a/src/rust/cxx/kj-rs-io/tests/http-test.c++ b/src/rust/cxx/kj-rs-io/tests/http-test.c++ new file mode 100644 index 00000000000..55dfb9506f9 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/tests/http-test.c++ @@ -0,0 +1,147 @@ +// kj-http integration ("the acid test"): a real kj::HttpServer serving on a kj-rs-io listener +// and a kj::HttpClient over a kj-rs-io connection, all driven by the tokio event loop. Proves +// that kj-http works unchanged over tokio-backed streams. + +#include "kj-rs-io/async-io.h" + +#include +#include +#include +#include + +#include + +namespace kj_rs_io_test { +namespace { + +using kj_rs_io::setupTokioAsyncIo; + +// Echoes the request body back as the response body, streaming (pumpTo), preserving the +// content length when known. +class EchoService final: public kj::HttpService { + public: + explicit EchoService(kj::HttpHeaderTable &table): table(table) {} + + kj::Promise request(kj::HttpMethod method, + kj::StringPtr url, + const kj::HttpHeaders &headers, + kj::AsyncInputStream &requestBody, + Response &response) override { + kj::HttpHeaders responseHeaders(table); + auto body = response.send(200, "OK", responseHeaders, requestBody.tryGetLength()); + co_await requestBody.pumpTo(*body); + } + + private: + kj::HttpHeaderTable &table; +}; + +kj::Array makeBody(size_t size) { + auto data = kj::heapArray(size); + for (size_t i = 0; i < size; i++) { + data[i] = static_cast(('A' + i / 8192 + i * 13) & 0xff); + } + return data; +} + +// Writes `data` in chunks to the request body stream, then closes it. +kj::Promise writeBody( + kj::Own body, kj::ArrayPtr data) { + constexpr size_t CHUNK = 128 * 1024; + size_t offset = 0; + while (offset < data.size()) { + size_t n = kj::min(CHUNK, data.size() - offset); + co_await body->write(data.slice(offset, offset + n)); + offset += n; + } + // Dropping `body` (coroutine frame teardown) finishes the request body. +} + +KJ_TEST("bodyless GET works (regression: kj-http never awaits its header-write " + "queue for bodyless requests, relying on KJ hot-promise write semantics)") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + kj::HttpHeaderTable table; + EchoService service(table); + kj::HttpServer server(io.getTimer(), table, service); + auto listener = io.getNetwork().parseAddress("127.0.0.1", 0).wait(ws)->listen(); + auto listenTask = server.listenHttp(*listener).eagerlyEvaluate(nullptr); + auto addr = io.getNetwork().parseAddress(kj::str("127.0.0.1:", listener->getPort())).wait(ws); + auto connection = addr->connect().wait(ws); + auto client = kj::newHttpClient(table, *connection); + kj::HttpHeaders headers(table); + auto request = client->request(kj::HttpMethod::GET, "/x", headers, static_cast(0)); + auto response = request.response.wait(ws); + KJ_EXPECT(response.statusCode == 200); + auto body = response.body->readAllBytes().wait(ws); + KJ_EXPECT(body.size() == 0); +} + +KJ_TEST("kj-http round trip with streaming bodies over kj-rs-io streams on the " + "tokio loop") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + kj::HttpHeaderTable table; + EchoService service(table); + kj::HttpServer server(io.getTimer(), table, service); + + // Server side: kj::HttpServer accepting from a kj-rs-io ConnectionReceiver. + auto listener = io.getNetwork().parseAddress("127.0.0.1", 0).wait(ws)->listen(); + auto listenTask = server.listenHttp(*listener).eagerlyEvaluate( + [](kj::Exception &&e) { KJ_FAIL_EXPECT("HTTP server failed", e); }); + + // Client side: kj::HttpClient over a kj-rs-io connection. + auto addr = io.getNetwork().parseAddress(kj::str("127.0.0.1:", listener->getPort())).wait(ws); + auto connection = addr->connect().wait(ws); + auto client = kj::newHttpClient(table, *connection); + + // Round trip 1: 4 MB POST with a streamed request body, echoed back and read while the + // request body is still being written (full-duplex streaming through the tokio loop). + { + constexpr size_t SIZE = 4 * 1024 * 1024; + auto data = makeBody(SIZE); + + kj::HttpHeaders headers(table); + auto request = + client->request(kj::HttpMethod::POST, "/echo", headers, static_cast(SIZE)); + + auto writeTask = writeBody(kj::mv(request.body), data).eagerlyEvaluate(nullptr); + auto response = request.response.wait(ws); + KJ_EXPECT(response.statusCode == 200); + KJ_EXPECT(KJ_ASSERT_NONNULL(response.body->tryGetLength()) == SIZE); + + auto echoed = response.body->readAllBytes(SIZE + 1).wait(ws); + KJ_ASSERT(echoed.size() == SIZE); + KJ_ASSERT(memcmp(echoed.begin(), data.begin(), SIZE) == 0); + writeTask.wait(ws); + } + + // Round trip 2 on the same connection (keep-alive): small GET, empty echoed body. + { + kj::HttpHeaders headers(table); + auto request = + client->request(kj::HttpMethod::GET, "/again", headers, static_cast(0)); + auto response = request.response.wait(ws); + KJ_EXPECT(response.statusCode == 200); + auto body = response.body->readAllBytes().wait(ws); + KJ_EXPECT(body.size() == 0); + } + + // Round trip 3: chunked request body (no expected size -> Transfer-Encoding: chunked). + { + kj::HttpHeaders headers(table); + auto request = client->request(kj::HttpMethod::POST, "/chunked", headers); + auto data = makeBody(64 * 1024); + auto writeTask = writeBody(kj::mv(request.body), data).eagerlyEvaluate(nullptr); + auto response = request.response.wait(ws); + KJ_EXPECT(response.statusCode == 200); + auto echoed = response.body->readAllBytes().wait(ws); + KJ_ASSERT(echoed.size() == data.size()); + KJ_ASSERT(memcmp(echoed.begin(), data.begin(), data.size()) == 0); + writeTask.wait(ws); + } +} + +} // namespace +} // namespace kj_rs_io_test diff --git a/src/rust/cxx/kj-rs-io/tests/lib.rs b/src/rust/cxx/kj-rs-io/tests/lib.rs new file mode 100644 index 00000000000..b62a0893e5a --- /dev/null +++ b/src/rust/cxx/kj-rs-io/tests/lib.rs @@ -0,0 +1,84 @@ +#![allow(clippy::unused_async)] +#![allow(clippy::missing_errors_doc)] +#![allow(clippy::must_use_candidate)] +#![allow(clippy::missing_panics_doc)] +#![allow(clippy::missing_safety_doc)] +#![allow(clippy::unnecessary_box_returns)] // cxx bridge functions return Box by contract + +mod serve_helpers; +mod test_helpers; + +use serve_helpers::ServeEchoSession; +use serve_helpers::start_serve_echo; +use serve_helpers::start_take_socket_echo; +use test_helpers::create_prebound_listener_fd; +use test_helpers::native_write_via_kj_unwrap; +use test_helpers::native_write_via_unwrap; + +#[cxx::bridge(namespace = "kj_rs_io_test")] +mod ffi { + /// A pre-bound, *blocking* std TCP listener handed to C++ as a raw fd (the `--socket-fd` + /// scenario for `wrapListenSocketFd`). + struct PreboundListener { + fd: i32, + port: u16, + } + + extern "Rust" { + /// Recovers the native tokio TcpStream from an unwrapped kj-rs-io stream Box and writes + /// `data` natively (tokio readiness API, no FFI-per-byte), then closes the connection. + async fn native_write_via_unwrap(stream: Box, data: Vec) -> Result<()>; + + /// Same, but starts from a `kj::AsyncIoStream&`: unwraps it from the Rust side via + /// `kj_rs_io::unwrap_kj_stream` (the API a native Rust server's glue will use), then writes + /// natively. Leaves the C++ wrapper hollow. + async unsafe fn native_write_via_kj_unwrap<'a>( + stream: Pin<&'a mut KjAsyncIoStream>, + data: &'a [u8], + ) -> Result<()>; + + /// Binds 127.0.0.1:0 with std (blocking mode, like an inherited `--socket-fd` listener) + /// and releases it as a raw fd owned by the caller. + fn create_prebound_listener_fd() -> Result; + + // --- serve_kj_stream (serve_helpers.rs) + + /// One echo server over a served kj stream: `start_serve_echo` picks the transport + /// path (unwrap fast path or duplex pump) via `kj_rs_io::serve_kj_stream` — taking + /// ownership of the stream — and spawns the echo consumer on the loop runtime; + /// `drive()` runs the connection (the pump, on the pumped path) to completion — + /// dropping the `drive()` promise mid-connection is the abort-on-drop path, which + /// also destroys the owned stream. + type ServeEchoSession; + + fn start_serve_echo(stream: KjOwn) -> Box; + + /// Like `start_serve_echo`, but through the native-only `take_kj_socket` entry point + /// (unwrap tier, else fd-dup tier): errors — instead of pumping — for streams that + /// are neither. The consumed stream is destroyed before this returns. + fn start_take_socket_echo(stream: KjOwn) -> Result>; + + /// Whether the unwrap fast path was taken (perf observability surface). + fn is_native(self: &ServeEchoSession) -> bool; + + /// Runs the connection to completion (see the type's docs). May only be called once. + async unsafe fn drive<'a>(self: &'a ServeEchoSession) -> Result<()>; + + /// Resolves once the echo task has exited — observing that dropping `drive()` EOFs + /// the pumped consumer. + async unsafe fn wait_echo_done<'a>(self: &'a ServeEchoSession); + } + + extern "Rust" { + #[namespace = "kj_rs_io"] + type TokioStream = kj_rs_io::TokioStream; + } + + unsafe extern "C++" { + include!("kj-rs-io/unwrap.h"); + + #[namespace = "kj"] + #[cxx_name = "AsyncIoStream"] + type KjAsyncIoStream = kj_rs_io::KjAsyncIoStream; + } +} diff --git a/src/rust/cxx/kj-rs-io/tests/serve-test.c++ b/src/rust/cxx/kj-rs-io/tests/serve-test.c++ new file mode 100644 index 00000000000..18629c784b7 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/tests/serve-test.c++ @@ -0,0 +1,284 @@ +// Tests for kj_rs_io::serve_kj_stream (serve.rs): the native-serve entry +// point Rust servers use to drive a kj::AsyncIoStream's connection natively. Following kj-rs +// conventions, C++ KJ_TESTs drive; Rust helpers (tests/serve_helpers.rs) run the echo server +// side over whichever transport path the entry point picks. + +#include "kj-rs-io-test/lib.rs.h" +#include "kj-rs-io/async-io.h" + +#include +#include +#include +#include +#include + +#include + +#if !_WIN32 +#include // getpid()/unlink() for the unix-socket pair helper +#endif + +namespace kj_rs_io_test { +namespace { + +using kj_rs_io::setupTokioAsyncIo; +using kj_rs_io::TokioAsyncIoContext; + +struct ConnectedPair { + kj::Own listener; + kj::Own client; + kj::Own server; +}; + +ConnectedPair makeTcpPair(TokioAsyncIoContext &io) { + auto &ws = io.getWaitScope(); + auto listener = io.getNetwork().parseAddress("127.0.0.1").wait(ws)->listen(); + auto connectAddr = + io.getNetwork().parseAddress(kj::str("127.0.0.1:", listener->getPort())).wait(ws); + auto acceptPromise = listener->accept(); + auto client = connectAddr->connect().wait(ws); + auto server = acceptPromise.wait(ws); + return ConnectedPair{kj::mv(listener), kj::mv(client), kj::mv(server)}; +} + +#if !_WIN32 +// A connected AF_UNIX stream-socket pair, from the kj-rs-io network's `unix:` support. Used to +// prove take_kj_socket serves a non-TCP fd. Binds a short /tmp path (unlinked before bind to +// clear stale sockets, and +// again once connected) unique per process + call, so parallel/repeated runs never collide. +ConnectedPair makeUnixPair(TokioAsyncIoContext &io) { + auto &ws = io.getWaitScope(); + static uint counter = 0; + auto path = kj::str("/tmp/kj-rs-io-serve-", ::getpid(), "-", counter++, ".sock"); + ::unlink(path.cStr()); + auto addr = kj::str("unix:", path); + auto listener = io.getNetwork().parseAddress(addr).wait(ws)->listen(); + auto connectAddr = io.getNetwork().parseAddress(addr).wait(ws); + auto acceptPromise = listener->accept(); + auto client = connectAddr->connect().wait(ws); + auto server = acceptPromise.wait(ws); + // The bound path is no longer needed once both ends are connected. + ::unlink(path.cStr()); + return ConnectedPair{kj::mv(listener), kj::mv(client), kj::mv(server)}; +} +#endif // !_WIN32 + +kj::Array makePatternedData(size_t size, kj::byte seed) { + auto data = kj::heapArray(size); + for (size_t i = 0; i < size; i++) { + data[i] = static_cast((i * 31 + seed) & 0xff); + } + return data; +} + +// The client side of an echo round trip: write `data` (in chunks) and concurrently read the +// echo back and verify it (concurrent, so bounded transports -- the pump duplex, socket +// buffers -- never deadlock on payloads larger than their buffering); then half-close and +// expect EOF. +kj::Promise echoRoundTrip( + kj::AsyncIoStream &clientStream, kj::ArrayPtr data) { + static auto constexpr readBack = [](kj::AsyncIoStream &s, + kj::ArrayPtr expected) -> kj::Promise { + auto buffer = kj::heapArray(expected.size()); + size_t total = co_await s.tryRead(buffer.begin(), buffer.size(), buffer.size()); + KJ_ASSERT(total == expected.size(), total, expected.size()); + KJ_ASSERT(memcmp(buffer.begin(), expected.begin(), expected.size()) == 0); + }; + + static auto constexpr writeAll = [](kj::AsyncIoStream &s, + kj::ArrayPtr data) -> kj::Promise { + constexpr size_t CHUNK = 64 * 1024; + size_t offset = 0; + while (offset < data.size()) { + size_t n = kj::min(CHUNK, data.size() - offset); + co_await s.write(data.slice(offset, offset + n)); + offset += n; + } + s.shutdownWrite(); // half-close: the echo server sees EOF and finishes flushing + }; + + co_await kj::joinPromisesFailFast( + kj::arr(writeAll(clientStream, data), readBack(clientStream, data))); + kj::byte extra; + KJ_EXPECT(co_await clientStream.tryRead(&extra, 1, 1) == 0); // EOF after echo completes +} + +// ======================================================================================= +// Unwrap fast path + +KJ_TEST("serve_kj_stream takes the native path for kj-rs-io TCP streams and " + "echoes") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pair = makeTcpPair(io); + + // Native path: the connection now belongs to the Rust side (the hollow wrapper was + // destroyed by serve_kj_stream). + auto session = start_serve_echo(kj::mv(pair.server)); + KJ_EXPECT(session->is_native()); + + auto data = makePatternedData(256 * 1024, 7); + auto drive = session->drive(); + auto client = echoRoundTrip(*pair.client, data); + kj::joinPromisesFailFast(kj::arr(kj::mv(drive), kj::mv(client))).wait(ws); +} + +// ======================================================================================= +// Duplex pump fallback (foreign streams) + +KJ_TEST("serve_kj_stream pumps foreign streams: bidirectional echo + half-close") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + // An in-memory kj pipe is the canonical foreign stream: not kj-rs-io-originated. + auto pipe = kj::newTwoWayPipe(); + + auto session = start_serve_echo(kj::mv(pipe.ends[0])); + KJ_EXPECT(!session->is_native()); + + auto data = makePatternedData(512 * 1024, 3); + auto drive = session->drive(); + auto client = echoRoundTrip(*pipe.ends[1], data); + // The pump owns `pipe.ends[0]` now; only the client end stays with the test. + kj::joinPromisesFailFast(kj::arr(kj::mv(drive), kj::mv(client))).wait(ws); +} + +KJ_TEST("serve_kj_stream pump: dropping the pump aborts the bridge (drop-abort)") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pipe = kj::newTwoWayPipe(); + + auto session = start_serve_echo(kj::mv(pipe.ends[0])); + KJ_EXPECT(!session->is_native()); + + { + // Prove the bridge is live: one small round trip, driving the pump only while the + // client operation runs, then *drop* the drive promise mid-connection. + auto drive = session->drive(); + auto oneRoundTrip = [](kj::AsyncIoStream &s) -> kj::Promise { + co_await s.write("ping"_kjb); + kj::byte buffer[4]; + size_t n = co_await s.tryRead(buffer, 4, 4); + KJ_ASSERT(n == 4); + KJ_ASSERT(memcmp(buffer, "ping", 4) == 0); + }(*pipe.ends[1]); + // exclusiveJoin: when the round trip finishes, `drive` is cancelled (dropped). + oneRoundTrip.exclusiveJoin(kj::mv(drive)).wait(ws); + } + + // Dropping the pump dropped the kj-side duplex end: the echo consumer reads EOF and + // finishes... + session->wait_echo_done().wait(ws); + + // ...and the pump destroyed the kj stream it owned: the peer observes teardown (a rejected + // write), not a zombie half-open pipe. + auto orphanWrite = kj::evalNow([&]() { return pipe.ends[1]->write("anyone there?"_kjb); }); + KJ_EXPECT(orphanWrite.poll(ws)); + orphanWrite + .then([]() { KJ_FAIL_EXPECT("write to a torn-down pipe unexpectedly succeeded"); }, + [](kj::Exception &&) { + }).wait(ws); +} + +// ======================================================================================= +// take_kj_socket (native-only: unwrap tier, else fd-dup tier) + +KJ_TEST("take_kj_socket unwraps kj-rs-io TCP streams (tier 1) and echoes") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pair = makeTcpPair(io); + + // Native socket taken; the (hollow) kj stream was destroyed inside take_kj_socket. + auto session = start_take_socket_echo(kj::mv(pair.server)); + KJ_EXPECT(session->is_native()); + + auto data = makePatternedData(256 * 1024, 5); + auto drive = session->drive(); + auto client = echoRoundTrip(*pair.client, data); + kj::joinPromisesFailFast(kj::arr(kj::mv(drive), kj::mv(client))).wait(ws); +} + +#if !_WIN32 +// A foreign kj::AsyncIoStream that exposes only its OS fd: unwrap must fail (it is not a +// kj-rs-io wrapper) and any per-read FFI would abort the test -- proving take_kj_socket's fd +// tier does all its I/O on the dup'd socket, never through the kj stream. +class FdOnlyStream final: public kj::AsyncIoStream { + public: + explicit FdOnlyStream(kj::Own inner): inner(kj::mv(inner)) {} + + kj::Maybe getFd() const override { + return inner->getFd(); + } + + kj::Promise tryRead(void *, size_t, size_t) override { + KJ_UNIMPLEMENTED("FdOnlyStream must not be read through the FFI"); + } + kj::Promise write(kj::ArrayPtr) override { + KJ_UNIMPLEMENTED("FdOnlyStream must not be written through the FFI"); + } + kj::Promise write(kj::ArrayPtr>) override { + KJ_UNIMPLEMENTED("FdOnlyStream must not be written through the FFI"); + } + kj::Promise whenWriteDisconnected() override { + KJ_UNIMPLEMENTED("FdOnlyStream must not be observed through the FFI"); + } + void shutdownWrite() override { + KJ_UNIMPLEMENTED("FdOnlyStream must not be shut down through the FFI"); + } + + private: + kj::Own inner; +}; + +KJ_TEST("take_kj_socket dups the fd of foreign fd-backed streams (tier 2); the " + "original stream may be dropped") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pair = makeTcpPair(io); + + // Hide the kj-rs-io origin behind a foreign wrapper: only getFd() is reachable. + auto foreign = kj::heap(kj::mv(pair.server)); + + // Fd tier: the dup is independent; take_kj_socket destroys the wrapper (and with it the + // original socket's owner) before returning, which must not tear the served connection down. + auto session = start_take_socket_echo(kj::mv(foreign)); + KJ_EXPECT(session->is_native()); + + auto data = makePatternedData(256 * 1024, 9); + auto drive = session->drive(); + auto client = echoRoundTrip(*pair.client, data); + kj::joinPromisesFailFast(kj::arr(kj::mv(drive), kj::mv(client))).wait(ws); +} + +// A unix-domain (AF_UNIX) socket must be served natively too: take_kj_socket's fd tier dups the +// fd and detects the family (getsockname), producing a tokio UnixStream (ServeIo::Unix). +KJ_TEST("take_kj_socket serves a unix-domain (AF_UNIX) socket via its fd tier and echoes") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto pair = makeUnixPair(io); + + // Hide the kj-rs-io origin behind a foreign wrapper so take_kj_socket must use its fd tier + // (dup + family detection), not the unwrap fast path -- proving the AF_UNIX -> UnixStream + // branch of serve_io_from_owned_fd. + auto foreign = kj::heap(kj::mv(pair.server)); + + auto session = start_take_socket_echo(kj::mv(foreign)); + KJ_EXPECT(session->is_native()); // ServeIo::Unix is a native path (no pump) + + auto data = makePatternedData(256 * 1024, 11); + auto drive = session->drive(); + auto client = echoRoundTrip(*pair.client, data); + kj::joinPromisesFailFast(kj::arr(kj::mv(drive), kj::mv(client))).wait(ws); +} +#endif // !_WIN32 + +KJ_TEST("take_kj_socket refuses fd-less foreign streams (no pump tier)") { + auto io = setupTokioAsyncIo(); + auto pipe = kj::newTwoWayPipe(); + + KJ_EXPECT_THROW_MESSAGE( + "cannot take the stream's socket natively", start_take_socket_echo(kj::mv(pipe.ends[0]))); +} + +} // namespace +} // namespace kj_rs_io_test diff --git a/src/rust/cxx/kj-rs-io/tests/serve_helpers.rs b/src/rust/cxx/kj-rs-io/tests/serve_helpers.rs new file mode 100644 index 00000000000..de1b13431ee --- /dev/null +++ b/src/rust/cxx/kj-rs-io/tests/serve_helpers.rs @@ -0,0 +1,109 @@ +//! Rust helpers for the `serve_kj_stream` `KJ_TEST`s (`serve-test.c++`): +//! echo sessions over each transport path, driven by the C++ tests. + +use std::cell::RefCell; + +use cxx::KjError; +use kj_rs::KjOwn; +use kj_rs_io::serve::ServeIo; +use kj_rs_io::serve::ServePath; +use kj_rs_io::serve::StreamPump; +use tokio::io::AsyncWriteExt; +use tokio::sync::watch; + +use crate::ffi::KjAsyncIoStream; + +type Result = std::result::Result; + +fn kj_err(message: impl std::fmt::Display) -> KjError { + KjError::new(cxx::KjExceptionType::Failed, message.to_string()) +} + +/// The echo task shared by every path: copy everything read back to the writer, then +/// propagate the half-close. Reports completion through `done_tx`. +async fn echo(io: ServeIo, done_tx: watch::Sender) -> std::io::Result { + let result = async { + let (mut rd, mut wr) = tokio::io::split(io); + let n = tokio::io::copy(&mut rd, &mut wr).await?; + wr.shutdown().await?; + Ok(n) + } + .await; + let _ = done_tx.send(true); + result +} + +/// One echo server over a served kj stream; see `start_serve_echo` in lib.rs. +pub struct ServeEchoSession { + native: bool, + /// Present for the pumped path; taken by `drive()`. + pump: RefCell>, + /// The echo task's completion signal (fires even if the task failed). + done_rx: watch::Receiver, + /// Taken by `drive()`. + echo: RefCell>>>, +} + +pub fn start_serve_echo(stream: KjOwn) -> Box { + let served = kj_rs_io::serve_kj_stream(stream); + Box::new(ServeEchoSession::new(served.io, served.pump)) +} + +/// Like `start_serve_echo`, but through the native-only entry point (`take_kj_socket`, +/// tiers 1 + 2): errors for streams that are neither kj-rs-io native nor fd-backed (the +/// handed-back stream is dropped with the error). +pub fn start_take_socket_echo(stream: KjOwn) -> Result> { + let io = kj_rs_io::take_kj_socket(stream).map_err(KjError::from)?; + Ok(Box::new(ServeEchoSession::new(io, None))) +} + +impl ServeEchoSession { + fn new(io: ServeIo, pump: Option) -> Self { + let native = io.path() == ServePath::Native; + let (done_tx, done_rx) = watch::channel(false); + // The echo consumer runs on this thread's KJ-loop runtime: the one-runtime shape + // (native sockets are registered with this runtime's I/O driver anyway). + let echo = kj_rs_tokio::spawn(echo(io, done_tx)); + Self { + native, + pump: RefCell::new(pump), + done_rx, + echo: RefCell::new(Some(echo)), + } + } + + pub fn is_native(&self) -> bool { + self.native + } + + /// Runs the connection to completion: drives the pump (if any) and waits for the echo + /// task to finish. Dropping the returned promise mid-connection drops the pump — + /// the abort-on-drop path. + pub async fn drive(&self) -> Result<()> { + let pump = self.pump.borrow_mut().take(); + let echo = self + .echo + .borrow_mut() + .take() + .ok_or_else(|| kj_err("drive() was already called"))?; + if let Some(pump) = pump { + pump.await?; + } + let echoed = echo + .await + .map_err(|e| kj_err(format!("echo task panicked: {e}")))? + .map_err(|e| kj_err(format!("echo failed: {e}")))?; + let _ = echoed; + Ok(()) + } + + /// Resolves once the echo task has finished (however `drive()` fared) — used by the + /// drop-abort test to observe that dropping the pump EOFs the consumer. + pub async fn wait_echo_done(&self) { + let mut rx = self.done_rx.clone(); + // Cannot fail: the sender is owned by the echo task, which always sends before exit; + // even if it panicked, the closed channel resolves wait_for with an error we ignore + // after checking the flag. + let _ = rx.wait_for(|done| *done).await; + } +} diff --git a/src/rust/cxx/kj-rs-io/tests/test_helpers.rs b/src/rust/cxx/kj-rs-io/tests/test_helpers.rs new file mode 100644 index 00000000000..b41b244dd86 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/tests/test_helpers.rs @@ -0,0 +1,70 @@ +use std::io; +use std::pin::Pin; + +use cxx::KjError; +use kj_rs_io::TokioStream; +use tokio::io::Interest; +use tokio::net::TcpStream; + +use crate::ffi::KjAsyncIoStream; +use crate::ffi::PreboundListener; + +type Result = std::result::Result; + +fn kj_err(message: impl std::fmt::Display) -> KjError { + KjError::new(cxx::KjExceptionType::Failed, message.to_string()) +} + +/// Write-all over the native tokio readiness API. +async fn native_write_all(stream: &TcpStream, data: &[u8]) -> io::Result<()> { + let mut written = 0; + while written < data.len() { + match stream.try_write(&data[written..]) { + Ok(n) => written += n, + Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + stream.ready(Interest::WRITABLE).await?; + } + Err(e) => return Err(e), + } + } + Ok(()) +} + +pub async fn native_write_via_unwrap(stream: Box, data: Vec) -> Result<()> { + let tcp = stream + .into_tcp_stream() + .ok_or_else(|| kj_err("expected a TCP stream"))?; + native_write_all(&tcp, &data).await.map_err(kj_err)?; + // Dropping `tcp` closes the connection; the C++ side observes data followed by EOF. + Ok(()) +} + +pub async fn native_write_via_kj_unwrap( + stream: Pin<&mut KjAsyncIoStream>, + data: &[u8], +) -> Result<()> { + // Safety: the C++ test guarantees no I/O promises are in flight on `stream`. + let native = unsafe { kj_rs_io::unwrap_kj_stream(stream) }.map_err(KjError::from)?; + let tcp = native + .into_tcp_stream() + .ok_or_else(|| kj_err("expected a TCP stream"))?; + native_write_all(&tcp, data).await.map_err(kj_err)?; + Ok(()) +} + +pub fn create_prebound_listener_fd() -> Result { + let listener = std::net::TcpListener::bind("127.0.0.1:0").map_err(kj_err)?; + let port = listener.local_addr().map_err(kj_err)?.port(); + #[cfg(unix)] + { + use std::os::fd::IntoRawFd; + Ok(PreboundListener { + fd: listener.into_raw_fd(), + port, + }) + } + #[cfg(not(unix))] + { + Err(kj_err("not supported on this platform")) + } +} diff --git a/src/rust/cxx/kj-rs-io/unwrap.h b/src/rust/cxx/kj-rs-io/unwrap.h new file mode 100644 index 00000000000..d9ce7109271 --- /dev/null +++ b/src/rust/cxx/kj-rs-io/unwrap.h @@ -0,0 +1,78 @@ +#pragma once +// Declarations needed by the kj-rs-io cxx bridge (lib.rs) itself. The full C++ API lives in +// kj-rs-io/async-io.h; this header only exposes what the generated bridge code references: +// kj::AsyncIoStream (as an opaque extern C++ type), the unwrap hook, and the bridged stream +// operations backing serve_kj_stream()'s pump fallback (serve.rs). + +#include + +#include + +namespace kj_rs_io { + +struct TokioStream; // Opaque Rust type, defined in the generated lib.rs.h. + +// Recovers the native Rust stream out of a kj::AsyncIoStream created by kj-rs-io (the "unwrap +// fast path"), leaving the wrapper hollow: any further I/O through the wrapper throws. Throws if +// `stream` is not a kj-rs-io stream or was already unwrapped, or if I/O promises are still in +// flight on it (caller contract; not detected). +// +// Implemented in async-io.c++. Rust code calls this through kj_rs_io::unwrap_kj_stream(). +::rust::Box unwrapTokioStream(kj::AsyncIoStream &stream); + +// --- Bridged operations on a *foreign* kj::AsyncIoStream (one that did not originate in +// kj-rs-io and therefore cannot be unwrapped). These back the duplex-pump fallback of +// serve_kj_stream() (serve.rs): the pump owns the stream and reads and writes it through +// concurrent Rust *shared* borrows — kj two-way streams support one read and one write in +// flight at once. + +// A Rust shared borrow (&KjAsyncIoStream) arrives in C++ as a const&: the constness is cxx's +// wire format for "shared", not a property of the object — the stream is uniquely owned by the +// pump and never actually const. NOTE the meaning mismatch with KJ convention: KJ const means +// *thread-safe* (Rust's Sync), but this shared-ness is the weaker property — reentrant use from +// one thread's async tasks (the bridge types are !Send/!Sync, so Rust can never move these +// borrows off the KJ event-loop thread that owns the stream). Recover the callable reference +// here, once. +inline kj::AsyncIoStream &pumpStream(const kj::AsyncIoStream &stream) { + return const_cast(stream); +} + +// Corresponds to kj::AsyncIoStream::tryRead(buffer, minBytes, buffer.size()). +inline kj::Promise kjStreamTryRead( + const kj::AsyncIoStream &stream, ::rust::Slice buffer, size_t minBytes) { + return pumpStream(stream).tryRead(buffer.data(), minBytes, buffer.size()); +} + +// Corresponds to kj::AsyncIoStream::write(buffer) (write-all semantics). +inline kj::Promise kjStreamWrite( + const kj::AsyncIoStream &stream, ::rust::Slice buffer) { + return pumpStream(stream).write(kj::arrayPtr(buffer.data(), buffer.size())); +} + +// Corresponds to kj::AsyncIoStream::shutdownWrite(). +inline void kjStreamShutdownWrite(const kj::AsyncIoStream &stream) { + pumpStream(stream).shutdownWrite(); +} + +// The stream's underlying raw OS socket handle -- a Unix fd (kj::AsyncIoStream::getFd()) or a +// win32 SOCKET (kj::AsyncIoStream::getWin32Handle()) -- widened to int64, or -1 if it exposes +// none. int64 fits both losslessly with one sentinel: a Unix fd is a non-negative int, and +// INVALID_SOCKET (~0 as UINT_PTR) is exactly -1 as int64 (live win64 SOCKET values fit in 32 +// bits per the Windows handle-interoperability guarantee, so they never collide with -1). +// Backs the handle tier of take_kj_socket() (ffi.rs); see that function's docs for the +// caller-asserted "the handle carries the stream's own bytes" contract (wrappers such as +// kj::TlsConnection forward getFd() to their *transport* socket, which this tier must never be +// used on). +inline int64_t kjStreamGetHandle(const kj::AsyncIoStream &stream) { +#if _WIN32 + // Validated by Windows CI; mirrors the unix arm. + KJ_IF_SOME(handle, stream.getWin32Handle()) { + return static_cast(reinterpret_cast(handle)); + } + return -1; +#else + return stream.getFd().orDefault(-1); +#endif +} + +} // namespace kj_rs_io diff --git a/src/rust/cxx/kj-rs-tokio/BUILD.bazel b/src/rust/cxx/kj-rs-tokio/BUILD.bazel new file mode 100644 index 00000000000..9e51af3e3f9 --- /dev/null +++ b/src/rust/cxx/kj-rs-tokio/BUILD.bazel @@ -0,0 +1,71 @@ +load("@rules_rust//rust:defs.bzl", "rust_library", "rust_test") +load("//:build/wd_cc_library.bzl", "wd_cc_library") +load("//src/rust/cxx/tools/bazel:rust_cxx_bridge.bzl", "rust_cxx_bridge") + +wd_cc_library( + name = "kj-rs-tokio-lib", + srcs = glob(["*.c++"]), + hdrs = glob(["*.h"]), + include_prefix = "kj-rs-tokio", + linkstatic = select({ + "@platforms//os:windows": True, + "//conditions:default": False, + }), + strip_include_prefix = "/src/rust/cxx/kj-rs-tokio", + visibility = ["//visibility:public"], + deps = [ + ":bridge", + # For kj-rs/waker.h + the kj_rs::futurePollArmNudge hook the port installs (the same-thread + # arm nudge). kj-rs is the abstract bridge core; no cycle (kj-rs does not depend on us). + "//src/rust/cxx/kj-rs:kj-rs-lib", + ], +) + +rust_library( + name = "kj-rs-tokio", + srcs = glob(["*.rs"]), + compile_data = glob(["*.h"]), + edition = "2024", + link_deps = [ + ":bridge", + ":kj-rs-tokio-lib", + ], + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + visibility = ["//visibility:public"], + deps = [ + "//src/rust/cxx", + "@crates_vendor//:tokio", + ], +) + +rust_test( + name = "kj-rs-tokio_test", + crate = "kj-rs-tokio", + edition = "2024", + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), +) + +rust_cxx_bridge( + name = "bridge", + src = "ffi.rs", + hdrs = glob(["*.h"]), + include_prefix = "kj-rs-tokio", + visibility = ["//visibility:public"], + deps = [ + "@capnp-cpp//src/kj:kj", + # kj-rs-tokio only needs the abstract async core (Promise / EventLoop / EventPort / + # Timer) to build the tokio-backed EventPort; it uses neither the abstract kj streams + # (:kj-async-io) nor the OS event loop (setupAsyncIo / UnixEventPort). Depending on + # :kj-async-core (not the :kj-async umbrella) keeps the tokio event loop off the + # concrete kj OS I/O layer (:kj-async-os), which the workerd rust-io hermeticity aspect + # forbids. + "@capnp-cpp//src/kj:kj-async-core", + "//src/rust/cxx:core", + ], +) diff --git a/src/rust/cxx/kj-rs-tokio/ffi.rs b/src/rust/cxx/kj-rs-tokio/ffi.rs new file mode 100644 index 00000000000..08b367b2acc --- /dev/null +++ b/src/rust/cxx/kj-rs-tokio/ffi.rs @@ -0,0 +1,209 @@ +//! The `#[cxx::bridge]` FFI island for kj-rs-tokio. +//! +//! This is the crate's single dedicated FFI-island file (file-top `#![allow(unsafe_code)]`). It +//! holds two kinds of hand-written `unsafe`: +//! +//! - the `#[cxx::bridge] mod bridge` — the C++ <-> Rust wire the C++ `kj_rs_tokio::TokioEventPort` +//! drives (see `tokio-event-port.h`); and +//! - the OS-specific precise absolute sleeps for the [`crate::port`] `HiResTimer` thread +//! (`boost_current_thread_priority` / `sleep_until`) — hand-rolled FFI (three raw syscall +//! symbols) rather than a `libc` dependency; all targets workerd builds for are 64-bit. +//! +//! Everything else — `lib.rs` and the entire event-port / timer business logic in `port.rs` — is +//! wholly-safe, compiler-proven unsafe-free under the crate-root `#![deny(unsafe_code)]`. +//! +//! [`crate::port`]: crate::port +#![allow(unsafe_code)] + +// Only the unix precise-sleep helpers below use `Instant`; the bridge itself is cross-platform. +#[cfg(unix)] +use std::time::Instant; + +use crate::port::TokioPort; +use crate::port::new_tokio_port; + +#[cxx::bridge(namespace = "kj_rs_tokio")] +// FFI island: the cxx bridge macro generates the `unsafe` extern shims. +// unnecessary_box_returns: returning the opaque `TokioPort` to C++ boxed is the cxx idiom. The +// lint's firing is platform-dependent (it has a size threshold and `TokioPort`'s size differs by +// target), so `#[expect]` would be unfulfilled on some targets. +#[expect(clippy::allow_attributes)] +#[allow(clippy::unnecessary_box_returns)] +mod bridge { + extern "Rust" { + type TokioPort; + + fn new_tokio_port() -> Box; + + /// Block until `wake()` is called or the KJ event loop becomes runnable, running tokio + /// tasks in the meantime. Returns the wake latch (see `TokioPort::take_wake_latch`). + fn wait_forever(&self) -> bool; + + /// Like `wait_forever`, but additionally returns after `timeout_ns` nanoseconds. The + /// C++ side computes the timeout from `kj::TimerImpl::timeoutToNextEvent()`. + fn wait_timeout_ns(&self, timeout_ns: u64) -> bool; + + /// Non-blocking: let the tokio scheduler run already-ready tasks for a bounded number of + /// turns. Never sleeps. Returns the wake latch. + fn poll(&self) -> bool; + + /// Set the wake latch and unblock a concurrent `wait_*`. Callable from any thread. + fn wake(&self); + + /// Called (on the loop thread only) when the KJ event loop becomes runnable. If the + /// thread is currently parked inside `wait_*`, unblock it so the KJ queue gets serviced. + fn notify_runnable(&self); + } +} + +/// Opts the timer thread out of OS timer-coalescing slop as far as an unprivileged process +/// can. On macOS, default-QoS threads get proportional timer leeway (~25-30% observed: +/// a 500 µs `mach_wait_until` overshoots by ~150 µs); `QOS_CLASS_USER_INTERACTIVE` shrinks +/// it substantially. Called once at timer-thread startup; best-effort. +#[cfg(target_os = "macos")] +pub fn boost_current_thread_priority() { + use core::ffi::c_int; + use core::ffi::c_uint; + const QOS_CLASS_USER_INTERACTIVE: c_uint = 0x21; + unsafe extern "C" { + fn pthread_set_qos_class_self_np(qos_class: c_uint, relative_priority: c_int) -> c_int; + } + // Safety: simple syscall wrapper acting on the calling thread; no memory crosses. + let _ = unsafe { pthread_set_qos_class_self_np(QOS_CLASS_USER_INTERACTIVE, 0) }; +} + +/// On Linux the equivalent knob is the per-thread hrtimer slack (default ~50 µs); shrink it +/// to 1 ns for this thread only. Best-effort. +#[cfg(target_os = "linux")] +pub fn boost_current_thread_priority() { + use core::ffi::c_int; + use core::ffi::c_ulong; + const PR_SET_TIMERSLACK: c_int = 29; + unsafe extern "C" { + fn prctl( + option: c_int, + arg2: c_ulong, + arg3: c_ulong, + arg4: c_ulong, + arg5: c_ulong, + ) -> c_int; + } + // Safety: simple syscall wrapper acting on the calling thread; no memory crosses. + let _ = unsafe { prctl(PR_SET_TIMERSLACK, 1, 0, 0, 0) }; +} + +#[cfg(all(unix, not(any(target_os = "macos", target_os = "linux"))))] +pub fn boost_current_thread_priority() {} + +/// Sleeps until `deadline` using `mach_wait_until`, which takes an *absolute* time in mach +/// tick units and honors it with microsecond-level precision (relative `nanosleep` on macOS +/// is subject to aggressive timer coalescing). +#[cfg(target_os = "macos")] +pub fn sleep_until(deadline: Instant) { + use core::ffi::c_int; + use std::sync::OnceLock; + + #[repr(C)] + struct MachTimebaseInfo { + numer: u32, + denom: u32, + } + unsafe extern "C" { + fn mach_absolute_time() -> u64; + fn mach_timebase_info(info: *mut MachTimebaseInfo) -> c_int; + fn mach_wait_until(deadline: u64) -> c_int; + } + static TIMEBASE: OnceLock<(u64, u64)> = OnceLock::new(); + + let Some(remaining) = deadline.checked_duration_since(Instant::now()) else { + return; + }; + let &(numer, denom) = TIMEBASE.get_or_init(|| { + let mut info = MachTimebaseInfo { numer: 0, denom: 0 }; + // Safety: `info` is a valid out-pointer for the duration of the call. + let rc = unsafe { mach_timebase_info(&raw mut info) }; + assert_eq!(rc, 0, "mach_timebase_info failed"); + (u64::from(info.numer), u64::from(info.denom)) + }); + // mach ticks -> ns is `ticks * numer / denom`, so ns -> ticks is `ns * denom / numer`. + let nanos = u64::try_from(remaining.as_nanos()).unwrap_or(u64::MAX); + let ticks = u64::try_from(u128::from(nanos) * u128::from(denom) / u128::from(numer)) + .unwrap_or(u64::MAX); + // Safety: no memory crosses the boundary; both calls are simple syscall wrappers. + unsafe { + let _ = mach_wait_until(mach_absolute_time().saturating_add(ticks)); + } +} + +/// Sleeps until `deadline` using `clock_nanosleep(CLOCK_MONOTONIC, TIMER_ABSTIME)`, which +/// is backed by hrtimers (microsecond-level precision, subject only to the default ~50 µs +/// timer slack). +#[cfg(target_os = "linux")] +pub fn sleep_until(deadline: Instant) { + use core::ffi::c_int; + use core::ffi::c_long; + + // Layout of glibc/musl `struct timespec` on the 64-bit targets workerd builds for. + #[repr(C)] + struct Timespec { + tv_sec: c_long, + tv_nsec: c_long, + } + const CLOCK_MONOTONIC: c_int = 1; + const TIMER_ABSTIME: c_int = 1; + const EINTR: c_int = 4; + unsafe extern "C" { + fn clock_gettime(clockid: c_int, tp: *mut Timespec) -> c_int; + fn clock_nanosleep( + clockid: c_int, + flags: c_int, + request: *const Timespec, + remain: *mut Timespec, + ) -> c_int; + } + + let Some(remaining) = deadline.checked_duration_since(Instant::now()) else { + return; + }; + let mut ts = Timespec { + tv_sec: 0, + tv_nsec: 0, + }; + // Safety: `ts` is a valid out-pointer for the duration of the call. + if unsafe { clock_gettime(CLOCK_MONOTONIC, &raw mut ts) } != 0 { + return; + } + ts.tv_sec = ts + .tv_sec + .saturating_add(c_long::try_from(remaining.as_secs()).unwrap_or(c_long::MAX)); + ts.tv_nsec += c_long::from(remaining.subsec_nanos()); + if ts.tv_nsec >= 1_000_000_000 { + ts.tv_sec += 1; + ts.tv_nsec -= 1_000_000_000; + } + loop { + // Safety: `ts` is a valid, initialized timespec; `remain` may be null with + // TIMER_ABSTIME (the absolute deadline makes restart-after-signal lossless). + let rc = unsafe { + clock_nanosleep( + CLOCK_MONOTONIC, + TIMER_ABSTIME, + &raw const ts, + std::ptr::null_mut(), + ) + }; + // clock_nanosleep returns the error number directly (not via errno). + if rc != EINTR { + break; + } + } +} + +/// Portable fallback for other unixes: plain relative sleep (coarser, but still typically +/// far better than a ~1 ms timer wheel). +#[cfg(all(unix, not(any(target_os = "macos", target_os = "linux"))))] +pub fn sleep_until(deadline: Instant) { + if let Some(remaining) = deadline.checked_duration_since(Instant::now()) { + std::thread::sleep(remaining); + } +} diff --git a/src/rust/cxx/kj-rs-tokio/lib.rs b/src/rust/cxx/kj-rs-tokio/lib.rs new file mode 100644 index 00000000000..224050d23ea --- /dev/null +++ b/src/rust/cxx/kj-rs-tokio/lib.rs @@ -0,0 +1,45 @@ +//! Rust half of the tokio-backed KJ event loop foundation. +//! +//! This crate owns a per-thread tokio `current_thread` runtime and exposes the primitives the +//! C++ `kj_rs_tokio::TokioEventPort` (see `tokio-event-port.h` for the full contract) needs to +//! implement `kj::EventPort`: parking the thread in `Runtime::block_on` (which drives the +//! whole tokio scheduler while C++ is "blocked"), a bounded non-blocking `poll`, and the +//! cross-thread `wake()` latch that `kj::Executor` and +//! `kj::newPromiseAndCrossThreadFulfiller` depend on. + +// Safety & panic enforcement walls. Test code exempted. +// +// `unsafe` is quarantined into a single named FFI island: the crate root denies `unsafe_code`, so +// the entire event-port business logic (`TokioPort`, the `wait`/`poll`/`wake` machinery, the +// `HiResTimer` orchestration) is *compiler-proven* to contain no hand-written unsafe. The one +// island that opts back in via `#![allow(unsafe_code)]` is `ffi.rs`: the `#[cxx::bridge]` wire plus +// the OS-syscall wrappers for precise short sleeps. +#![deny(unsafe_op_in_unsafe_fn)] +#![deny(unsafe_code)] +#![deny(clippy::undocumented_unsafe_blocks)] +#![deny( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::unreachable, + clippy::todo, + clippy::unimplemented +)] +#![cfg_attr( + test, + allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::unreachable, + clippy::todo, + clippy::unimplemented + ) +)] + +pub use port::TokioPort; +pub use port::current_handle; +pub use port::spawn; + +mod ffi; +mod port; diff --git a/src/rust/cxx/kj-rs-tokio/port.rs b/src/rust/cxx/kj-rs-tokio/port.rs new file mode 100644 index 00000000000..c05572b1feb --- /dev/null +++ b/src/rust/cxx/kj-rs-tokio/port.rs @@ -0,0 +1,653 @@ +//! Per-thread tokio `current_thread` runtime management and the Rust half of `TokioEventPort`. + +use std::cell::RefCell; +use std::future::Future; +use std::rc::Rc; +use std::sync::Arc; +#[cfg(unix)] +use std::sync::Condvar; +#[cfg(unix)] +use std::sync::Mutex; +#[cfg(unix)] +use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; +#[cfg(unix)] +use std::time::Instant; + +use tokio::runtime::Builder; +use tokio::runtime::Handle; +use tokio::runtime::Runtime; +use tokio::sync::Notify; +use tokio::task::JoinHandle; +use tokio::task::LocalSet; + +#[cfg(unix)] +use crate::ffi; + +/// How many scheduler turns `poll()` grants the runtime. Each `yield_now` re-queues the main +/// future at the back of the run queue, so every already-ready spawned task gets a chance to run +/// (repeatedly, up to the budget) without ever parking the thread. +/// +/// The value is a latency/throughput compromise, not derived from any tokio internal: large +/// enough to drain a typical burst of already-ready tasks in one `poll()` call, small enough to +/// bound how long `poll()` withholds control from the KJ loop when spawned tasks keep re-readying +/// each other. Safe to retune if profiling shows either starvation or excessive poll latency. +const POLL_YIELD_BUDGET: u32 = 16; + +/// Timeouts strictly below this go through the [`HiResTimer`] short-sleep path: tokio's timer +/// wheel has ~1 ms granularity, which would quantize sub-millisecond KJ timers (e.g. a 100 µs +/// `timer.afterDelay`) to a ~1 ms sleep. At and above a couple of milliseconds the wheel's +/// error is proportionally small, so long sleeps stay on the plain tokio path. +#[cfg(unix)] +const HIRES_TIMEOUT_THRESHOLD: Duration = Duration::from_millis(2); + +thread_local! { + /// Handle to the `TokioPort` runtime driving the KJ event loop on this thread, if any. + /// Registered by `TokioPort::new` and cleared on drop. Used by code (e.g. kj-rs-io) that + /// needs to *enter* the runtime context so tokio resources can register with its I/O driver + /// and timers. + static LOOP_RUNTIME_HANDLE: RefCell> = const { RefCell::new(None) }; + + /// The `LocalSet` onto which [`spawn`] enqueues tasks, and which `wait_*`/`poll` drive. Held + /// behind an `Rc` so `spawn` (which has no `&TokioPort`) can reach it while the port keeps + /// driving it. The `LocalSet` is `!Send`/`!Sync` and lives entirely on the loop thread, which + /// is why `TokioPort` itself does NOT hold it (it must stay `Send + Sync` for cross-thread + /// `wake()`); ownership lives here and is dropped in `TokioPort::drop`. + static LOOP_LOCAL_SET: RefCell>> = const { RefCell::new(None) }; +} + +/// Returns a handle to this thread's KJ-loop tokio runtime, if a `TokioEventPort` exists on this +/// thread. +#[must_use] +pub fn current_handle() -> Option { + LOOP_RUNTIME_HANDLE.with(|h| h.borrow().clone()) +} + +/// Returns this thread's KJ-loop `LocalSet`, if a `TokioEventPort` exists on this thread. +fn current_local_set() -> Option> { + LOOP_LOCAL_SET.with(|l| l.borrow().clone()) +} + +/// Spawns a future onto this thread's KJ-loop tokio runtime. +/// +/// The task runs whenever the KJ event loop sleeps (i.e. whenever C++ is blocked in +/// `promise.wait(waitScope)` or pumping via `poll()`), driven by the port's [`LocalSet`]. +/// +/// The future is spawned with [`LocalSet::spawn_local`], so it is pinned to this (the loop) +/// thread and does **not** need to be `Send`: the per-thread `current_thread` runtime never +/// migrates a task to another thread. This is what lets bridged futures that hold `!Send` KJ +/// handles (`OwnPromiseNode`, `kj::Own`, ...) be spawned directly. `JoinHandle`/drop-cancels +/// semantics are the same as `tokio::spawn`. +/// +/// # Panics +/// +/// Panics if no `TokioEventPort` has been created on this thread. +pub fn spawn(future: F) -> JoinHandle +where + F: Future + 'static, + F::Output: 'static, +{ + #[expect( + clippy::expect_used, + reason = "documented `# Panics` on this public API: calling spawn() without first creating a TokioEventPort on this thread is a caller-contract violation, not a recoverable runtime path" + )] + let local = current_local_set() + .expect("no kj-rs-tokio runtime on this thread; create a TokioEventPort first"); + local.spawn_local(future) +} + +/// State shared with `wake()` callers on other threads. +struct SharedState { + /// Unblocks the `block_on(...)` inside `wait_*` when `wake()` or `notify_runnable()` fires. + notify: Notify, + + /// The `kj::EventPort::wake()` latch: set by `wake()`, consumed (swapped to `false`) by the + /// return value of `wait_*`/`poll`. The KJ event loop uses a `true` return to know it must + /// drain cross-thread events (`kj::Executor`, `CrossThreadPromiseFulfiller`). + woken: AtomicBool, + + /// True while the loop thread is parked inside `wait_*`'s `block_on`. Only mutated from the + /// loop thread; read by `notify_runnable` (also loop-thread-only, but kept atomic so the + /// whole struct is `Sync` and the flag is safe if that ever changes). + sleeping: AtomicBool, +} + +/// A lazily-started dedicated thread delivering high-resolution wakeups for short +/// `wait_timeout_ns` sleeps (see [`HIRES_TIMEOUT_THRESHOLD`]). +/// +/// For a short timeout, `wait_impl` arms this timer with an absolute deadline before parking in +/// `block_on`; the thread performs a precise absolute sleep (`mach_wait_until` on macOS, +/// `clock_nanosleep(TIMER_ABSTIME)` on Linux) and then pokes the port's `Notify`. This composes +/// with every other wake-up source because it *is* the same mechanism (`wake()` and +/// `notify_runnable()` hit the same `Notify`); the tokio-side `timeout` stays armed as a coarse +/// backstop, so a lost or late high-res wakeup only degrades to the wheel's ~1 ms behavior, +/// never a hang. While unused the thread parks on a condvar (zero CPU); joined on drop. +#[cfg(unix)] +struct HiResTimer { + shared: Arc, + /// Lazily spawned by the first `arm()`; joined on drop. Only the loop thread arms, but the + /// handle sits behind a mutex so the struct is `Sync` without further reasoning. + thread: Mutex>>, +} + +#[cfg(unix)] +struct HiResShared { + /// The port state whose `notify` the timer thread pokes when a deadline is reached. + port: Arc, + request: Mutex, + condvar: Condvar, +} + +#[cfg(unix)] +struct HiResRequest { + /// Absolute deadline of the currently armed request, if any. + deadline: Option, + /// Bumped by `arm()` and `disarm()`. A wakeup is only delivered if the generation still + /// matches once the sleep finishes, so a canceled request (the wait was woken early by + /// `wake()`/`notify_runnable()`) does not leak a stale `Notify` permit into a later wait. + /// A lost race here is harmless: spurious early returns are explicitly allowed by the + /// `kj::EventPort::wait()` contract. + generation: u64, + shutdown: bool, +} + +#[cfg(unix)] +impl HiResTimer { + fn new(port: Arc) -> Self { + Self { + shared: Arc::new(HiResShared { + port, + request: Mutex::new(HiResRequest { + deadline: None, + generation: 0, + shutdown: false, + }), + condvar: Condvar::new(), + }), + thread: Mutex::new(None), + } + } + + /// Requests a `notify` on the port at `deadline`. Called on the loop thread right before it + /// parks in `block_on`. + fn arm(&self, deadline: Instant) { + { + let mut req = self + .shared + .request + .lock() + .unwrap_or_else(PoisonError::into_inner); + req.deadline = Some(deadline); + req.generation += 1; + } + self.shared.condvar.notify_one(); + // Lazily start the thread on first use, so loops that never schedule sub-millisecond + // timers never pay for it. + let mut thread = self.thread.lock().unwrap_or_else(PoisonError::into_inner); + if thread.is_none() { + let shared = Arc::clone(&self.shared); + #[expect( + clippy::expect_used, + reason = "OS thread spawn only fails under resource exhaustion; the hi-res timer thread is required for sub-millisecond KJ timer precision and there is no recovery at this site (fail-fast). Graceful degradation to tokio-wheel-only timing is tracked as a design debt." + )] + let thread_handle = std::thread::Builder::new() + .name("kj-rs-hires-timer".into()) + .spawn(move || hires_timer_main(&shared)) + .expect("failed to spawn kj-rs hires timer thread"); + *thread = Some(thread_handle); + } + } + + /// Cancels any armed request. An in-flight precise sleep cannot be interrupted, but the + /// generation bump turns its delivery into a no-op; worst case it delays the *next* arm's + /// wakeup by up to the threshold, which the tokio backstop bounds. + fn disarm(&self) { + let mut req = self + .shared + .request + .lock() + .unwrap_or_else(PoisonError::into_inner); + req.deadline = None; + req.generation += 1; + } +} + +#[cfg(unix)] +impl Drop for HiResTimer { + fn drop(&mut self) { + { + let mut req = self + .shared + .request + .lock() + .unwrap_or_else(PoisonError::into_inner); + req.shutdown = true; + } + self.shared.condvar.notify_one(); + let handle = self + .thread + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take(); + if let Some(handle) = handle { + let _ = handle.join(); + } + } +} + +#[cfg(unix)] +fn hires_timer_main(shared: &HiResShared) { + ffi::boost_current_thread_priority(); + let mut req = shared + .request + .lock() + .unwrap_or_else(PoisonError::into_inner); + loop { + if req.shutdown { + return; + } + if let Some(deadline) = req.deadline.take() { + let generation = req.generation; + drop(req); + ffi::sleep_until(deadline); + req = shared + .request + .lock() + .unwrap_or_else(PoisonError::into_inner); + if req.shutdown { + return; + } + if req.generation == generation { + // Still the request we were armed with: unblock the loop thread. `Notify` + // stores a permit if the loop has not reached `notified()` yet, so this wakeup + // cannot be lost. + shared.port.notify.notify_one(); + } + } else { + // Nothing armed: park until the next arm() or shutdown. Zero CPU while idle. + req = shared + .condvar + .wait(req) + .unwrap_or_else(PoisonError::into_inner); + } + } +} + +/// The Rust backing of one `kj_rs_tokio::TokioEventPort` (C++). Owns the per-thread +/// `current_thread` tokio runtime. +/// +/// One instance per KJ event loop, created on (and driven by) that loop's thread. Only `wake()` +/// may be called from other threads. +pub struct TokioPort { + runtime: Runtime, + state: Arc, + /// High-resolution wakeup source for sub-millisecond `wait_timeout_ns` sleeps; see + /// [`HIRES_TIMEOUT_THRESHOLD`]. On non-unix targets short sleeps stay on the tokio wheel. + #[cfg(unix)] + hires: HiResTimer, +} + +// `wake()` is called through a `&TokioPort` shared with arbitrary threads, so the type must be +// `Sync` (`Runtime`, `Notify` and the atomics all are). +const _: () = { + const fn assert_send_sync() {} + assert_send_sync::(); +}; + +// Opaque cxx types must cross the bridge boxed. The lint's firing is platform-dependent (it has +// a size threshold and `TokioPort`'s size differs by target), so `#[expect]` would be unfulfilled +// on some targets. +#[expect(clippy::allow_attributes)] +#[allow(clippy::unnecessary_box_returns)] +pub fn new_tokio_port() -> Box { + Box::new(TokioPort::new()) +} + +impl TokioPort { + /// # Panics + /// + /// Panics if the tokio runtime cannot be built. + #[must_use] + pub fn new() -> Self { + #[expect( + clippy::expect_used, + reason = "startup-only: building the per-thread current_thread runtime fails only under resource exhaustion, at which point fail-fast at port construction is the correct behavior" + )] + let runtime = Builder::new_current_thread() + .enable_time() + // The I/O driver dispatches readiness for tokio sockets (used by kj-rs-io's + // tokio-backed KJ streams). It is only *driven* while this runtime is inside + // `block_on` (i.e. in `wait_*`/`poll`), which is exactly when the KJ loop sleeps. + .enable_io() + .build() + .expect("failed to build current_thread tokio runtime"); + LOOP_RUNTIME_HANDLE.with(|h| { + let mut slot = h.borrow_mut(); + assert!( + slot.is_none(), + "a kj-rs-tokio runtime already exists on this thread (one KJ event loop per \ + thread, hence one TokioEventPort per thread)" + ); + *slot = Some(runtime.handle().clone()); + }); + // The `LocalSet` that `spawn()` enqueues onto and `wait_*`/`poll` drive. Owned by the + // thread-local (not by `TokioPort`, which must stay `Send + Sync`); dropped in `drop()`. + LOOP_LOCAL_SET.with(|l| { + *l.borrow_mut() = Some(Rc::new(LocalSet::new())); + }); + let state = Arc::new(SharedState { + notify: Notify::new(), + woken: AtomicBool::new(false), + sleeping: AtomicBool::new(false), + }); + Self { + runtime, + #[cfg(unix)] + hires: HiResTimer::new(Arc::clone(&state)), + state, + } + } + + /// Handle to this port's runtime, usable to spawn tasks from any thread. + #[must_use] + pub fn handle(&self) -> Handle { + self.runtime.handle().clone() + } + + pub(crate) fn wait_forever(&self) -> bool { + self.wait_impl(None) + } + + pub(crate) fn wait_timeout_ns(&self, timeout_ns: u64) -> bool { + self.wait_impl(Some(Duration::from_nanos(timeout_ns))) + } + + fn wait_impl(&self, timeout: Option) -> bool { + let state = &self.state; + state.sleeping.store(true, Ordering::SeqCst); + + // Sub-millisecond deadlines cannot be met by tokio's ~1 ms timer wheel: hand them to + // the high-resolution timer thread (see HiResTimer). + #[cfg(unix)] + let hires_armed = match timeout { + Some(t) if !t.is_zero() && t < HIRES_TIMEOUT_THRESHOLD => { + self.hires.arm(Instant::now() + t); + true + } + _ => false, + }; + + // This `block_on` is where tokio owns the thread: it drives *all* tasks — those spawned + // onto the port's `LocalSet` via `spawn()` (driven by `LocalSet::block_on`'s `run_until`) + // as well as any `tokio::spawn`ed tasks on the current_thread runtime — not just the + // future passed to it. Wake-up sources: `wake()` from another thread (with the `woken` + // latch set), `notify_runnable()` (a task inside this very `block_on` re-entered C++ and + // armed a KJ event), and the next KJ timer deadline — via the tokio wheel for long + // sleeps, via the high-res timer thread (same `notify`) for sub-millisecond ones. + // `Notify` stores a permit if `notify_one()` arrives before `notified()` is polled, so + // there is no lost-wakeup window; spurious early returns are explicitly allowed by the + // `kj::EventPort::wait()` contract. + #[expect( + clippy::expect_used, + reason = "TokioPort::new registers this thread's LocalSet; wait() only runs while this port drives its own thread, so the LocalSet is always present — absence is an unreachable internal invariant" + )] + let local = + current_local_set().expect("TokioPort is driving without a registered LocalSet"); + local.block_on(&self.runtime, async { + match timeout { + Some(t) => { + let _ = tokio::time::timeout(t, state.notify.notified()).await; + } + None => state.notify.notified().await, + } + }); + + #[cfg(unix)] + if hires_armed { + self.hires.disarm(); + } + + state.sleeping.store(false, Ordering::SeqCst); + self.take_wake_latch() + } + + pub(crate) fn poll(&self) -> bool { + // Bounded, non-blocking pump: the main future is always immediately ready to run again + // after `yield_now`, so the scheduler never parks; it just interleaves any ready spawned + // tasks (LocalSet + runtime) with our yields until the budget is spent. + #[expect( + clippy::expect_used, + reason = "TokioPort::new registers this thread's LocalSet; poll() only runs while this port drives its own thread, so the LocalSet is always present — absence is an unreachable internal invariant" + )] + let local = + current_local_set().expect("TokioPort is driving without a registered LocalSet"); + local.block_on(&self.runtime, async { + for _ in 0..POLL_YIELD_BUDGET { + tokio::task::yield_now().await; + } + }); + self.take_wake_latch() + } + + pub(crate) fn wake(&self) { + self.state.woken.store(true, Ordering::SeqCst); + self.state.notify.notify_one(); + } + + pub(crate) fn notify_runnable(&self) { + // Only meaningful while parked in `wait_impl`; `setRunnable(true)` is only ever called on + // the loop thread, so if `sleeping` is set we are inside `block_on` and the caller is a + // tokio task that re-entered C++ and armed a KJ event. Without this nudge the loop would + // keep sleeping until the next timer/wake even though it has work queued. + if self.state.sleeping.load(Ordering::SeqCst) { + self.state.notify.notify_one(); + } + } + + /// Consumes the wake latch: returns `true` iff `wake()` was called since the last `true` + /// return from `wait_*`/`poll`. + fn take_wake_latch(&self) -> bool { + self.state.woken.swap(false, Ordering::SeqCst) + } +} + +impl Default for TokioPort { + fn default() -> Self { + Self::new() + } +} + +impl Drop for TokioPort { + fn drop(&mut self) { + // `try_with`, not `with`: a port owned by an object destroyed during thread/process + // teardown (e.g. under KJ_CLEAN_SHUTDOWN) can be dropped after this thread's Rust TLS + // destructors have already run, where `with` panics with an AccessError. If the TLS is + // gone, its values (including the LocalSet) were already dropped with it, so there is + // nothing left to clear. + let _ = LOOP_RUNTIME_HANDLE.try_with(|h| { + *h.borrow_mut() = None; + }); + // Drop the LocalSet, canceling all still-pending spawned tasks. This runs on the loop + // thread (the only thread that may drop a `!Send` LocalSet), after + // `kj::WaitScope`/`kj::EventLoop` destruction (see `TokioAsyncIoContext` member order), + // so canceled tasks must not touch KJ objects from their `Drop` impls. + let _ = LOOP_LOCAL_SET.try_with(|l| { + *l.borrow_mut() = None; + }); + // Dropping `self.runtime` cancels any tasks spawned via `tokio::spawn` (as opposed to the + // LocalSet). Same threading/ordering constraints as above. + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn wake_latch_semantics() { + let port = TokioPort::new(); + // No wake: a timed-out wait reports false. + assert!(!port.wait_timeout_ns(1_000_000)); + // Wake before wait: latch is reported exactly once. + port.wake(); + assert!(port.wait_timeout_ns(1_000_000)); + assert!(!port.wait_timeout_ns(1_000_000)); + // Wake is also consumed by poll(). + port.wake(); + assert!(port.poll()); + assert!(!port.poll()); + } + + #[test] + fn wake_from_other_thread_unblocks_wait_forever() { + let port = Arc::new(TokioPort::new()); + let port2 = Arc::clone(&port); + let thread = std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(10)); + port2.wake(); + }); + assert!(port.wait_forever()); + thread.join().unwrap(); + } + + #[test] + fn spawned_tasks_run_during_wait() { + let port = TokioPort::new(); + let (tx, mut rx) = tokio::sync::oneshot::channel::(); + let mut jh = spawn(async move { + tokio::time::sleep(Duration::from_millis(5)).await; + tx.send(42).unwrap(); + }); + // The task only runs inside wait_impl's block_on. + let mut done = false; + for _ in 0..100 { + let _ = port.wait_timeout_ns(20_000_000); + if let Ok(v) = rx.try_recv() { + assert_eq!(v, 42); + done = true; + break; + } + } + assert!(done); + // The JoinHandle should complete promptly now. + port.runtime.block_on(&mut jh).unwrap(); + } + + #[test] + fn poll_never_sleeps() { + let port = TokioPort::new(); + // A pending spawned task must not make poll() block. + let _jh = spawn(std::future::pending::<()>()); + let start = std::time::Instant::now(); + assert!(!port.poll()); + assert!(start.elapsed() < Duration::from_millis(100)); + } + + /// `spawn` accepts `!Send` futures because it is backed by `LocalSet::spawn_local`. This + /// future holds an `Rc` — which is `!Send` — across an await point, exercising that path, + /// and proves such a task actually runs to completion on the loop thread. + #[test] + fn spawn_accepts_non_send_futures() { + use std::cell::Cell; + let port = TokioPort::new(); + let counter = Rc::new(Cell::new(0u32)); + let task_counter = Rc::clone(&counter); + // Detached on purpose; the `Rc` capture makes the future `!Send`. + let _jh = spawn(async move { + for _ in 0..3 { + tokio::task::yield_now().await; + } + task_counter.set(task_counter.get() + 1); + }); + let mut done = false; + for _ in 0..100 { + let _ = port.wait_timeout_ns(1_000_000); + if counter.get() == 1 { + done = true; + break; + } + } + assert!(done, "non-Send spawned task did not run to completion"); + } + + /// The high-res short-sleep path: a 100 µs timeout must not be quantized to tokio's ~1 ms + /// timer wheel. + /// + /// The wheel failure mode is a FLOOR — under quantization every sample takes >= ~1 ms — + /// while CI load only inflates some samples (loaded macOS CI VMs push the MEDIAN past + /// 500 µs). So assert on the minimum of 31 samples: load can't push all of them up, the + /// wheel floor pushes every one of them past the bound. + #[cfg(unix)] + #[test] + fn short_timeouts_are_sub_millisecond() { + let port = TokioPort::new(); + // Warm-up: lazily spawns the hires timer thread and faults in the block_on paths. + let _ = port.wait_timeout_ns(100_000); + + let mut samples: Vec = (0..31) + .map(|_| { + let start = std::time::Instant::now(); + let _ = port.wait_timeout_ns(100_000); + start.elapsed() + }) + .collect(); + samples.sort(); + let fastest = samples[0]; + assert!( + fastest >= Duration::from_micros(100), + "woke before the deadline: fastest {fastest:?}" + ); + assert!( + fastest < Duration::from_micros(500), + "100us timeout quantized: fastest {fastest:?}" + ); + } + + /// Long sleeps stay on the plain tokio path and remain accurate (and, by construction, + /// never touch the hires thread — see `HIRES_TIMEOUT_THRESHOLD`). + #[test] + fn long_timeouts_still_accurate() { + let port = TokioPort::new(); + let start = std::time::Instant::now(); + let _ = port.wait_timeout_ns(20_000_000); + let elapsed = start.elapsed(); + assert!( + elapsed >= Duration::from_millis(19), + "woke early: {elapsed:?}" + ); + assert!( + elapsed < Duration::from_millis(500), + "woke far too late: {elapsed:?}" + ); + } + + /// `wake()` must still interrupt a wait that has the high-res timer armed, and the stale + /// hires wakeup for the canceled deadline must not corrupt the latch of later waits. + #[cfg(unix)] + #[test] + fn wake_interrupts_short_timeout_wait() { + let port = Arc::new(TokioPort::new()); + // Arm the hires machinery once so the thread exists. + let _ = port.wait_timeout_ns(100_000); + + let port2 = Arc::clone(&port); + let thread = std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(2)); + port2.wake(); + }); + // A chain of short waits; one of them must observe the wake latch. + let mut woken = false; + for _ in 0..1000 { + if port.wait_timeout_ns(500_000) { + woken = true; + break; + } + } + assert!(woken); + thread.join().unwrap(); + // The latch was consumed; subsequent short waits time out normally with latch false. + assert!(!port.wait_timeout_ns(100_000)); + } +} diff --git a/src/rust/cxx/kj-rs-tokio/tests/BUILD.bazel b/src/rust/cxx/kj-rs-tokio/tests/BUILD.bazel new file mode 100644 index 00000000000..8e47a751db2 --- /dev/null +++ b/src/rust/cxx/kj-rs-tokio/tests/BUILD.bazel @@ -0,0 +1,52 @@ +load("@rules_cc//cc:cc_test.bzl", "cc_test") +load("@rules_rust//rust:defs.bzl", "rust_library") +load("//src/rust/cxx/tools/bazel:rust_cxx_bridge.bzl", "rust_cxx_bridge") + +rust_library( + name = "tests", + srcs = glob(["*.rs"]), + edition = "2024", + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + deps = [ + ":bridge", + "//src/rust/cxx", + "//src/rust/cxx/kj-rs", + "//src/rust/cxx/kj-rs-tokio", + "@crates_vendor//:tokio", + ], +) + +rust_cxx_bridge( + name = "bridge", + src = "lib.rs", + include_prefix = "kj-rs-tokio-test", + deps = [ + "//src/rust/cxx/kj-rs", + ], +) + +cc_test( + name = "tokio-event-port-test", + size = "medium", + srcs = [ + "tokio-event-port-test.c++", + ], + linkstatic = select({ + "@platforms//os:windows": True, + "//conditions:default": False, + }), + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + deps = [ + ":bridge", + ":tests", + "//src/rust/cxx/kj-rs-tokio:kj-rs-tokio-lib", + "//src/rust/cxx/third-party:runtime", + "@capnp-cpp//src/kj:kj-test", + ], +) diff --git a/src/rust/cxx/kj-rs-tokio/tests/lib.rs b/src/rust/cxx/kj-rs-tokio/tests/lib.rs new file mode 100644 index 00000000000..d136096062f --- /dev/null +++ b/src/rust/cxx/kj-rs-tokio/tests/lib.rs @@ -0,0 +1,47 @@ +#![allow(clippy::unused_async)] +#![allow(clippy::missing_errors_doc)] +#![allow(clippy::must_use_candidate)] +#![allow(clippy::missing_panics_doc)] + +mod test_helpers; + +use test_helpers::completed_task_count; +use test_helpers::has_loop_runtime_handle; +use test_helpers::spawn_pending_task; +use test_helpers::spawn_task_on_runtime; +use test_helpers::threaded_wake_future; +use test_helpers::tokio_sleep_on_runtime; + +type Result = std::io::Result; +type Error = std::io::Error; + +#[cxx::bridge(namespace = "kj_rs_tokio_test")] +mod ffi { + extern "Rust" { + /// Spawns a task on this thread's KJ-loop tokio runtime which sleeps `delay_ms` (real + /// tokio time) and then produces `value`; awaits its `JoinHandle`. The returned future is + /// bridged to a `kj::Promise` by kj-rs and polled by the KJ event loop; the spawned task + /// itself only runs while the loop sleeps inside `TokioEventPort::wait()`/`poll()`. + async fn spawn_task_on_runtime(delay_ms: u64, value: u32) -> Result; + + /// Spawns a `tokio::time::sleep` on the loop runtime and awaits it (via the task's + /// `JoinHandle`). + async fn tokio_sleep_on_runtime(delay_ms: u64) -> Result<()>; + + /// A bridged Rust future that, on its first poll, spawns a task on the loop runtime which + /// sleeps ~10ms (real tokio time) and then wakes a clone of the future's waker. The wake + /// therefore arrives same-thread (the runtime is driven by the port on the loop thread), + /// exercising the future⇄promise bridge's asynchronous same-thread re-drive path. + async fn threaded_wake_future() -> Result<()>; + + /// Detaches a never-completing task onto the loop runtime (for teardown testing: the + /// runtime must drop it cleanly when the port is destroyed). + fn spawn_pending_task(); + + /// Number of `spawn_task_on_runtime` tasks that ran to completion (process-wide). + fn completed_task_count() -> u64; + + /// True if `kj_rs_tokio::current_handle()` finds a runtime on this thread. + fn has_loop_runtime_handle() -> bool; + } +} diff --git a/src/rust/cxx/kj-rs-tokio/tests/test_helpers.rs b/src/rust/cxx/kj-rs-tokio/tests/test_helpers.rs new file mode 100644 index 00000000000..27a2a12a3a0 --- /dev/null +++ b/src/rust/cxx/kj-rs-tokio/tests/test_helpers.rs @@ -0,0 +1,89 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::time::Duration; + +use crate::Error; +use crate::Result; + +/// Process-wide counter of completed `spawn_task_on_runtime` tasks, so C++ tests can verify the +/// spawned task really ran (on the loop runtime) rather than the value being produced some other +/// way. +static COMPLETED_TASKS: AtomicU64 = AtomicU64::new(0); + +pub fn completed_task_count() -> u64 { + COMPLETED_TASKS.load(Ordering::SeqCst) +} + +pub fn has_loop_runtime_handle() -> bool { + kj_rs_tokio::current_handle().is_some() +} + +pub async fn spawn_task_on_runtime(delay_ms: u64, value: u32) -> Result { + let join_handle = kj_rs_tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(delay_ms)).await; + COMPLETED_TASKS.fetch_add(1, Ordering::SeqCst); + value + }); + join_handle.await.map_err(Error::other) +} + +pub async fn tokio_sleep_on_runtime(delay_ms: u64) -> Result<()> { + let join_handle = kj_rs_tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(delay_ms)).await; + }); + join_handle.await.map_err(Error::other) +} + +pub fn spawn_pending_task() { + // Deliberately detached: dropping the JoinHandle does not cancel the task; it stays pending + // in the runtime until the runtime itself is dropped with the TokioEventPort. + drop(kj_rs_tokio::spawn(std::future::pending::<()>())); +} + +/// A future that, on its first poll, spawns a task on the loop's own tokio runtime which sleeps +/// briefly and then wakes a clone of the future's waker. Because that runtime is driven by the +/// `TokioEventPort` on the event loop's own thread, the wake arrives *same-thread*. Exercises a +/// bridged future being suspended and re-driven by an asynchronous same-thread wake (via a cloned +/// waker) under the tokio-backed event port. +struct DelayedWakeFuture { + spawned: bool, + done: Arc, +} + +impl Future for DelayedWakeFuture { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { + if self.done.load(Ordering::SeqCst) { + return Poll::Ready(()); + } + if !self.spawned { + self.spawned = true; + let waker = cx.waker().clone(); + let done = Arc::clone(&self.done); + // Detached task on the loop runtime; it runs on the loop thread when the port drives + // the runtime, so `waker.wake()` arms the FuturePollEvent same-thread. + drop(kj_rs_tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(10)).await; + done.store(true, Ordering::SeqCst); + waker.wake(); + })); + } + Poll::Pending + } +} + +pub async fn threaded_wake_future() -> Result<()> { + DelayedWakeFuture { + spawned: false, + done: Arc::new(AtomicBool::new(false)), + } + .await; + Ok(()) +} diff --git a/src/rust/cxx/kj-rs-tokio/tests/tokio-event-port-test.c++ b/src/rust/cxx/kj-rs-tokio/tests/tokio-event-port-test.c++ new file mode 100644 index 00000000000..780da6e8ba8 --- /dev/null +++ b/src/rust/cxx/kj-rs-tokio/tests/tokio-event-port-test.c++ @@ -0,0 +1,428 @@ +// Tests for TokioEventPort / setupTokioAsyncIo: a kj::EventLoop driven by a per-thread tokio +// current_thread runtime. Following kj-rs conventions, C++ KJ_TESTs drive; Rust helpers (see +// tests/lib.rs, bridged by workerd-cxx) provide async behaviors. + +#include "kj-rs-tokio-test/lib.rs.h" +#include "kj-rs-tokio/tokio-event-port.h" + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace kj_rs_tokio_test { +namespace { + +using kj_rs_tokio::setupTokioAsyncIo; + +void delayMillis(uint64_t millis) { + std::this_thread::sleep_for(std::chrono::milliseconds(millis)); +} + +// ======================================================================================= +// Basics: the KJ event loop works as usual on top of the tokio-backed port. + +KJ_TEST("promises resolve on a TokioEventPort loop") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + KJ_EXPECT(kj::evalLater([]() { return 123; }).wait(ws) == 123); + KJ_EXPECT(kj::Promise(42).wait(ws) == 42); +} + +KJ_TEST("evalLater ordering is preserved") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + kj::Vector order; + // eagerlyEvaluate arms each promise immediately (KJ promises are otherwise lazy: continuations + // only get scheduled once awaited), so all three events sit in the queue in creation order. + auto p1 = kj::evalLater([&]() { order.add(1); }).eagerlyEvaluate(nullptr); + auto p2 = kj::evalLater([&]() { order.add(2); }).eagerlyEvaluate(nullptr); + auto p3 = kj::evalLater([&]() { order.add(3); }).eagerlyEvaluate(nullptr); + + // Waiting on the last promise runs all three events in FIFO order. + p3.wait(ws); + p1.wait(ws); + p2.wait(ws); + + KJ_ASSERT(order.size() == 3); + KJ_EXPECT(order[0] == 1); + KJ_EXPECT(order[1] == 2); + KJ_EXPECT(order[2] == 3); +} + +KJ_TEST("promise chains resolve") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + auto promise = + kj::evalLater([]() { return 1; }).then([](int v) { + return kj::Promise(v + 1); + }).then([](int v) { return v * 10; }); + KJ_EXPECT(promise.wait(ws) == 20); +} + +KJ_TEST("evalLast fires when the loop would sleep") { + // kj::evalLast events live on the would-sleep queue, which is only serviced through the + // EventLoop::poll() path (EventLoop::wait() switches to poll() when would-sleep waiters + // exist). This test hangs or misorders if the port's poll() is broken. + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + kj::Vector order; + auto last = kj::evalLast([&]() { order.add(2); }); + auto later = kj::evalLater([&]() { order.add(1); }); + + later.wait(ws); + // evalLast must not have run yet: the loop never ran out of work while waiting on `later`. + KJ_ASSERT(order.size() == 1); + KJ_EXPECT(order[0] == 1); + + last.wait(ws); + KJ_ASSERT(order.size() == 2); + KJ_EXPECT(order[1] == 2); +} + +// ======================================================================================= +// Timers: kj::TimerImpl fed by the port; advanceTo() after every wait()/poll(). + +KJ_TEST("timer.afterDelay fires with real elapsed time") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto &timer = io.getTimer(); + + auto &sysClock = kj::systemPreciseMonotonicClock(); + auto before = sysClock.now(); + auto timerBefore = timer.now(); + + // If the port forgot timerImpl.advanceTo() after waits, this would never resolve (caught by + // the test timeout). + timer.afterDelay(30 * kj::MILLISECONDS).wait(ws); + + KJ_EXPECT(sysClock.now() - before >= 30 * kj::MILLISECONDS); + // Timer time is synced to the monotonic clock at each wait return. + KJ_EXPECT(timer.now() - timerBefore >= 30 * kj::MILLISECONDS); +} + +KJ_TEST("multiple timers fire in deadline order") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto &timer = io.getTimer(); + + kj::Vector order; + auto p3 = timer.afterDelay(30 * kj::MILLISECONDS).then([&]() { + order.add(3); + }).eagerlyEvaluate(nullptr); + auto p1 = + timer.afterDelay(5 * kj::MILLISECONDS).then([&]() { order.add(1); }).eagerlyEvaluate(nullptr); + auto p2 = timer.afterDelay(15 * kj::MILLISECONDS).then([&]() { + order.add(2); + }).eagerlyEvaluate(nullptr); + + p3.wait(ws); + KJ_ASSERT(order.size() == 3); + KJ_EXPECT(order[0] == 1); + KJ_EXPECT(order[1] == 2); + KJ_EXPECT(order[2] == 3); + p1.wait(ws); + p2.wait(ws); +} + +KJ_TEST("timer fires while blocked waiting on a cross-thread event") { + // The port must bound each sleep by timeoutToNextEvent(): the loop first wakes at the timer + // deadline (long before the cross-thread fulfill), fires the timer, then goes back to sleep. + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto &timer = io.getTimer(); + + auto paf = kj::newPromiseAndCrossThreadFulfiller(); + bool timerFired = false; + auto timerPromise = timer.afterDelay(10 * kj::MILLISECONDS).then([&]() { + timerFired = true; + }).eagerlyEvaluate(nullptr); + + kj::Thread thread([fulfiller = kj::mv(paf.fulfiller)]() mutable { + delayMillis(100); + fulfiller->fulfill(); + }); + + paf.promise.wait(ws); + KJ_EXPECT(timerFired); + timerPromise.wait(ws); +} + +#if !_WIN32 +// ======================================================================================= +// Sub-millisecond timer precision: the high-resolution short-sleep path (unix-only: +// mach_wait_until on macOS, clock_nanosleep(TIMER_ABSTIME) on Linux). Without it, tokio's +// ~1 ms timer wheel quantizes a 100 µs kj::Timer delay to ~1.1 ms (a ~10x slowdown observed in +// workerd's Timed100us stream-piping benchmarks). Windows stays on the wheel, so these bounds +// don't apply there. + +KJ_TEST("sub-millisecond afterDelay is not quantized to tokio's ~1ms timer wheel") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto &timer = io.getTimer(); + auto &sysClock = kj::systemPreciseMonotonicClock(); + + // Warm-up: the first short wait lazily spawns the high-res timer thread. + timer.afterDelay(100 * kj::MICROSECONDS).wait(ws); + + auto sampleDelays = [&](kj::Duration delay) { + kj::Vector samples; + for (int i = 0; i < 31; i++) { + auto timerBefore = timer.now(); + auto before = sysClock.now(); + timer.afterDelay(delay).wait(ws); + samples.add(sysClock.now() - before); + // Never-early contract, checked on the timer's own clock: afterDelay computes its deadline + // from timer.now(), which only advances when the loop wakes, so it lags the wall clock + // slightly -- wall-clock elapsed can legitimately undershoot `delay` by that staleness + // (observed as a 492 us "500 us" sleep under ASan). + KJ_EXPECT(timer.now() - timerBefore >= delay, timer.now() - timerBefore, delay); + } + std::sort(samples.begin(), samples.end()); + return samples; + }; + + // Bound rationale: the wheel-quantized failure mode is a FLOOR -- tokio rounds every + // sub-millisecond delay up to the next ~1 ms tick, so under quantization no sample can complete + // in less than ~1 ms. The high-res path completes a 100 µs sleep in ~115-150 µs on idle + // hardware; loaded CI runners (macOS VMs especially) inflate the median well past 500 µs, but + // even there some of 31 samples land near the ideal. So assert on the MINIMUM sample: load + // noise can't push all 31 samples up, while the wheel floor pushes every one of them >= ~1 ms. + auto samples100 = sampleDelays(100 * kj::MICROSECONDS); + KJ_EXPECT(samples100[0] < 500 * kj::MICROSECONDS, samples100[0] / kj::MICROSECONDS); + + auto samples500 = sampleDelays(500 * kj::MICROSECONDS); + KJ_EXPECT(samples500[0] < 1000 * kj::MICROSECONDS, samples500[0] / kj::MICROSECONDS); +} + +KJ_TEST("sequential 100us timers run at ~100us each, not ~1ms each") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto &timer = io.getTimer(); + auto &sysClock = kj::systemPreciseMonotonicClock(); + + timer.afterDelay(100 * kj::MICROSECONDS).wait(ws); // warm-up + + constexpr int kIterations = 200; + kj::Duration minIteration = kj::maxValue; + auto timerBefore = timer.now(); + for (int i = 0; i < kIterations; i++) { + auto iterBefore = sysClock.now(); + timer.afterDelay(100 * kj::MICROSECONDS).wait(ws); + minIteration = kj::min(minIteration, sysClock.now() - iterBefore); + } + // Lower bound on the timer's own clock (see the sampling test above for why wall-clock elapsed + // can undershoot). + auto timerElapsed = timer.now() - timerBefore; + + // Ideal total is 20 ms; measured is ~25-40 ms idle but >140 ms on loaded macOS CI VMs, so a + // total-time bound can't separate load from the wheel-quantized failure mode (>= ~220 ms). + // Quantization is a floor, though: under it EVERY iteration takes >= ~1 ms, while under mere + // load at least one of 200 iterations still lands near the ~100 µs ideal. Assert on the + // fastest iteration. + KJ_EXPECT(timerElapsed >= kIterations * 100 * kj::MICROSECONDS, timerElapsed / kj::MILLISECONDS); + KJ_EXPECT(minIteration < 500 * kj::MICROSECONDS, minIteration / kj::MICROSECONDS); +} + +KJ_TEST("wake() from another thread interrupts short-timer waits promptly") { + // While the loop is continuously in the high-res short-wait path (a re-chaining 200 µs + // timer), a cross-thread fulfill must still get through promptly: the high-res wakeup and + // wake() share the same Notify, and the wake latch must survive interleaved timer wakeups. + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto &timer = io.getTimer(); + auto &sysClock = kj::systemPreciseMonotonicClock(); + + struct Chain { + static kj::Promise run(kj::Timer &timer) { + return timer.afterDelay(200 * kj::MICROSECONDS).then([&timer]() { return run(timer); }); + } + }; + auto keepBusy = Chain::run(timer).eagerlyEvaluate(nullptr); + + auto paf = kj::newPromiseAndCrossThreadFulfiller(); + auto before = sysClock.now(); + kj::Thread thread([fulfiller = kj::mv(paf.fulfiller)]() mutable { + delayMillis(5); + fulfiller->fulfill(7); + }); + + KJ_EXPECT(paf.promise.wait(ws) == 7); + auto elapsed = sysClock.now() - before; + KJ_EXPECT(elapsed >= 5 * kj::MILLISECONDS, elapsed / kj::MILLISECONDS); + KJ_EXPECT(elapsed < 100 * kj::MILLISECONDS, elapsed / kj::MILLISECONDS); +} + +KJ_TEST("long sleeps remain accurate alongside the high-res short-sleep path") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto &timer = io.getTimer(); + auto &sysClock = kj::systemPreciseMonotonicClock(); + + // Prime the high-res thread, then take a long sleep: it must go through the plain tokio + // path (no spin, no early wake from stale short-timer state). + timer.afterDelay(100 * kj::MICROSECONDS).wait(ws); + + auto before = sysClock.now(); + timer.afterDelay(20 * kj::MILLISECONDS).wait(ws); + auto elapsed = sysClock.now() - before; + KJ_EXPECT(elapsed >= 20 * kj::MILLISECONDS, elapsed / kj::MILLISECONDS); + KJ_EXPECT(elapsed < 500 * kj::MILLISECONDS, elapsed / kj::MILLISECONDS); +} +#endif // !_WIN32 + +// ======================================================================================= +// Cross-thread: the wake() -> wait()-returns-true latch is what drains kj::Executor events and +// cross-thread fulfillers. + +KJ_TEST("executeAsync from another thread runs on the tokio-ported loop") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + const kj::Executor &executor = kj::getCurrentThreadExecutor(); + auto paf = kj::newPromiseAndFulfiller(); + auto fulfiller = kj::mv(paf.fulfiller); + + kj::Thread thread([&executor, &fulfiller]() { + // A plain portless loop so this thread can wait on the cross-thread promise. + kj::EventLoop loop; + kj::WaitScope threadWs(loop); + kj::uint result = executor + .executeAsync([&fulfiller]() { + // Runs on the main (tokio-ported) loop. + fulfiller->fulfill(42); + return 99u; + }).wait(threadWs); + KJ_ASSERT(result == 99); + }); + + // While we are blocked here, the other thread's executeAsync must wake the port (wake() -> + // wait() returns true -> executor drained). + KJ_EXPECT(paf.promise.wait(ws) == 42); +} + +KJ_TEST("cross-thread fulfiller wakes a blocked wait()") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + auto paf = kj::newPromiseAndCrossThreadFulfiller(); + kj::Thread thread([fulfiller = kj::mv(paf.fulfiller)]() mutable { + delayMillis(10); // Give the main thread time to actually block in wait(). + fulfiller->fulfill(123); + }); + + KJ_EXPECT(paf.promise.wait(ws) == 123); +} + +// ======================================================================================= +// Rust integration: spawned tokio tasks run while C++ is blocked in promise.wait(), and the +// existing kj-rs future<->promise bridge works unchanged under the new port. + +KJ_TEST("Rust task spawned on the loop's runtime completes while C++ is " + "blocked in wait()") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + KJ_EXPECT(has_loop_runtime_handle()); + + uint64_t completedBefore = completed_task_count(); + auto promise = spawn_task_on_runtime(20, 42); + // This top-level wait() parks the thread inside the tokio runtime's block_on, which is what + // drives the spawned task (and its tokio sleep) to completion. + KJ_EXPECT(promise.wait(ws) == 42); + KJ_EXPECT(completed_task_count() == completedBefore + 1); +} + +KJ_TEST("tokio timers inside spawned tasks work") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + auto &sysClock = kj::systemPreciseMonotonicClock(); + auto before = sysClock.now(); + tokio_sleep_on_runtime(25).wait(ws); + KJ_EXPECT(sysClock.now() - before >= 25 * kj::MILLISECONDS); +} + +KJ_TEST("promise.poll() pumps the tokio scheduler without blocking") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + auto promise = spawn_task_on_runtime(30, 7); + // Not done yet; poll() must not sleep the 30ms away. + auto &sysClock = kj::systemPreciseMonotonicClock(); + auto before = sysClock.now(); + KJ_EXPECT(!promise.poll(ws)); + KJ_EXPECT(sysClock.now() - before < 30 * kj::MILLISECONDS); + + KJ_EXPECT(promise.wait(ws) == 7); +} + +KJ_TEST("kj-rs bridged Rust future with same-thread delayed waker works under the new " + "port") { + // A bridged Rust future that suspends, then is re-driven by a wake delivered from a task on the + // loop's own tokio runtime (same thread, via the TokioEventPort). Exercises the future⇄promise + // bridge's asynchronous same-thread re-drive path under the new port. + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + + threaded_wake_future().wait(ws); + + []() -> kj::Promise { co_await threaded_wake_future(); }().wait(ws); +} + +KJ_TEST("KJ coroutine can co_await spawned Rust tasks and KJ timers together") { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto &timer = io.getTimer(); + + int result = [&timer]() -> kj::Promise { + co_await timer.afterDelay(5 * kj::MILLISECONDS); + uint32_t value = co_await spawn_task_on_runtime(5, 11); + co_await timer.afterDelay(5 * kj::MILLISECONDS); + co_return static_cast(value) + 1; + }().wait(ws); + KJ_EXPECT(result == 12); +} + +// ======================================================================================= +// Teardown. + +KJ_TEST("context destruction with pending spawned tasks and armed timers is " + "clean") { + { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + auto &timer = io.getTimer(); + + // A detached, never-completing Rust task: must be dropped by the runtime at teardown. + spawn_pending_task(); + + // Armed timer and an in-flight spawned task; their promises are destroyed (canceling the + // KJ side) before the context itself, per declaration order. + auto timerPromise = timer.afterDelay(60 * kj::SECONDS); + auto spawnedPromise = spawn_task_on_runtime(60'000, 1); + KJ_EXPECT(!spawnedPromise.poll(ws)); + KJ_EXPECT(!timerPromise.poll(ws)); + } + + // The thread is fully cleaned up: a fresh context on the same thread works. + { + auto io = setupTokioAsyncIo(); + auto &ws = io.getWaitScope(); + KJ_EXPECT(kj::evalLater([]() { return 5; }).wait(ws) == 5); + KJ_EXPECT(has_loop_runtime_handle()); + } + KJ_EXPECT(!has_loop_runtime_handle()); +} + +} // namespace +} // namespace kj_rs_tokio_test diff --git a/src/rust/cxx/kj-rs-tokio/tokio-event-port.c++ b/src/rust/cxx/kj-rs-tokio/tokio-event-port.c++ new file mode 100644 index 00000000000..b23e14cb528 --- /dev/null +++ b/src/rust/cxx/kj-rs-tokio/tokio-event-port.c++ @@ -0,0 +1,90 @@ +#include "kj-rs-tokio/tokio-event-port.h" + +#include "kj-rs/waker.h" + +#include + +namespace kj_rs_tokio { + +namespace { +// The port active on this loop thread, used by the arm-nudge thunk installed into +// kj_rs::futurePollArmNudge. One port per thread (KJ's one-loop-per-thread model). +thread_local TokioEventPort* activePort = nullptr; + +// Captureless thunk (matches the `void(*)()` hook type) that nudges this thread's active port. +void portArmNudgeThunk() { + if (activePort != nullptr) { + activePort->nudge(); + } +} +} // namespace + +TokioEventPort::TokioEventPort() + : clock(kj::systemPreciseMonotonicClock()), + rustPort(new_tokio_port()), + timerImpl(clock.now()) { + // Install the same-thread arm-nudge hook so a tokio task that arms a KJ event during our + // block_on() park reliably wakes us even when KJ's edge-triggered setRunnable() misses it. + activePort = this; + kj_rs::futurePollArmNudge = &portArmNudgeThunk; +} + +TokioEventPort::~TokioEventPort() { + if (activePort == this) { + activePort = nullptr; + kj_rs::futurePollArmNudge = nullptr; + } +} + +bool TokioEventPort::wait() { + bool woken; + // Bound the sleep by the next KJ timer deadline, if any. `timeoutToNextEvent()` rounds up, so + // we always sleep until just *after* the timer is due. + KJ_IF_SOME(timeoutNs, timerImpl.timeoutToNextEvent(clock.now(), kj::NANOSECONDS, kj::maxValue)) { + woken = rustPort->wait_timeout_ns(timeoutNs); + } else { + woken = rustPort->wait_forever(); + } + + // Load-bearing: TimerImpl only fires timer events from advanceTo(). Forgetting this after a + // wait means every kj::Timer promise silently never resolves. + timerImpl.advanceTo(clock.now()); + return woken; +} + +bool TokioEventPort::poll() { + bool woken = rustPort->poll(); + timerImpl.advanceTo(clock.now()); + return woken; +} + +void TokioEventPort::wake() const { + // Callable from any thread; the Rust side latches the flag and unblocks a concurrent wait(). + rustPort->wake(); +} + +void TokioEventPort::setRunnable(bool runnable) { + // Called by the EventLoop (always on the loop's own thread) on empty<->runnable transitions. + // See the member comment on `runnable` for the future scheduled-pump hook point. + this->runnable = runnable; + if (runnable) { + // If we are currently parked inside wait()'s block_on, a tokio task just re-entered C++ and + // armed a KJ event; unblock so the loop can service its queue. No-op otherwise. + rustPort->notify_runnable(); + } +} + +void TokioEventPort::nudge() { + // Same-thread nudge: if we are parked inside wait()'s block_on, unblock so the loop services the + // KJ event a tokio task just armed. notify_runnable() no-ops when not parked. + rustPort->notify_runnable(); +} + +TokioAsyncIoContext setupTokioAsyncIo() { + auto port = kj::heap(); + auto loop = kj::heap(*port); + auto waitScope = kj::heap(*loop); + return TokioAsyncIoContext{kj::mv(port), kj::mv(loop), kj::mv(waitScope)}; +} + +} // namespace kj_rs_tokio diff --git a/src/rust/cxx/kj-rs-tokio/tokio-event-port.h b/src/rust/cxx/kj-rs-tokio/tokio-event-port.h new file mode 100644 index 00000000000..4ecb2f52df2 --- /dev/null +++ b/src/rust/cxx/kj-rs-tokio/tokio-event-port.h @@ -0,0 +1,109 @@ +#pragma once +// TokioEventPort: a kj::EventPort backed by a per-thread tokio current_thread runtime. +// Ownership inversion, not replacement: C++ keeps creating and awaiting kj::Promises exactly +// as today; what changes is *who sleeps*. When the KJ loop would block in the OS, wait() +// parks the thread inside `tokio::Runtime::block_on`, driving the whole tokio scheduler, so +// every Rust task on the loop's runtime makes progress while C++ is "blocked". +// +// Notes on the kj::EventPort contract (see kj/async.h): +// - wait()/poll() return true iff wake() latched. This is load-bearing: kj::Executor's +// executeAsync and kj::newPromiseAndCrossThreadFulfiller only get drained on `true`. +// - The kj::TimerImpl is advanced after every wait()/poll(). Sub-millisecond deadlines +// bypass tokio's ~1ms timer wheel via a dedicated high-resolution timer thread (see +// HiResTimer in port.rs). +// - One TokioEventPort (and hence one runtime) per thread, matching KJ's one-loop-per-thread +// model. The port, loop, and WaitScope must live and die on that thread. +// - What a tokio task on this runtime may do with KJ, today: it may complete a *bridged* +// future — make ready anything whose readiness reaches the loop through a Rust waker (a +// kj_rs::FutureWakerCell clone), e.g. send on a channel/oneshot that a bridged future is +// awaiting. Those wakes escape a parked wait() via the arm-nudge hook (see nudge()). It must +// NOT arm KJ events by any other means while the loop may be parked: re-entering C++ to +// fulfill a PromiseFulfiller / add to a TaskSet / resolve a promise chain, creating or +// awaiting a bridged (eager-by-default) kj promise, or arming a KJ timer all arm events that +// KJ's edge-filtered setRunnable() does not report during a park (lastRunnableState is stale- +// true), and nothing else touches the port's parked block_on — the arm sits unserviced until +// an unrelated wakeup (under wait-forever: possibly forever). Generalizing the nudge so tasks +// can use KJ promises freely is future work; until then, task->KJ communication must be +// waker-mediated only. Tasks must also never re-enter `promise.wait()` / +// `waitScope.poll()` on this thread: that nests block_on inside block_on, which tokio +// rejects (the panic surfaces as a kj::Exception). + +#include "kj-rs-tokio/ffi.rs.h" + +#include +#include +#include + +namespace kj_rs_tokio { + +class TokioEventPort final: public kj::EventPort { + public: + TokioEventPort(); + ~TokioEventPort(); + KJ_DISALLOW_COPY_AND_MOVE(TokioEventPort); + + // kj::EventPort implementation. + bool wait() override; + bool poll() override; + void wake() const override; + void setRunnable(bool runnable) override; + + // Nudge this port's Rust half out of a blocking wait() when a same-thread FutureWakerCell arm + // happens during block_on() that KJ's edge-triggered setRunnable() would miss (installed as + // kj_rs::futurePollArmNudge for this thread while the port lives). Idempotent: notify_runnable() + // no-ops unless the loop is parked. + void nudge(); + + // Timer fed by this port. now() is frozen while KJ events run and advances only when the loop + // waits/polls, preserving stock KJ timer semantics (the port calls timerImpl.advanceTo() after + // every wait()/poll() return; timers would silently never fire otherwise). + kj::Timer &getTimer() { + return timerImpl; + } + + // The Rust half (runtime + wake state). Rust code on this thread can also reach the runtime + // via kj_rs_tokio::current_handle() / kj_rs_tokio::spawn(). + const TokioPort &getRustPort() const { + return *rustPort; + } + + private: + const kj::MonotonicClock &clock; + ::rust::Box rustPort; + kj::TimerImpl timerImpl; + + // Recorded runnable state, updated by setRunnable() on empty<->runnable transitions. Today + // its only job is to nudge a concurrent wait() out of its block_on (a tokio task may have + // re-entered C++ and armed a KJ event while the loop slept); it is also the hook point for a + // future scheduled-pump model where Rust owns the thread and setRunnable(true) schedules a + // bounded waitScope.poll() pump task. + bool runnable = false; +}; + +// Mirrors the shape of kj::setupAsyncIo() (see kj/async-io.h) for the tokio-backed loop. Owns +// the event port (and thus the per-thread tokio runtime), the kj::EventLoop constructed with +// that port, and the kj::WaitScope. +struct TokioAsyncIoContext { + // Destroyed in reverse declaration order: waitScope, then loop (asserts its queue is empty), + // then port (dropping the tokio runtime, which cancels still-pending spawned tasks). Keep it + // this way: the loop must still exist while promises/timers are cancelled, and the runtime + // must outlive the loop because bridged Rust futures are cancelled through KJ promise + // destruction during loop teardown. + kj::Own port; + kj::Own loop; + kj::Own waitScope; + + TokioEventPort &getPort() { + return *port; + } + kj::WaitScope &getWaitScope() { + return *waitScope; + } + kj::Timer &getTimer() { + return port->getTimer(); + } +}; + +TokioAsyncIoContext setupTokioAsyncIo(); + +} // namespace kj_rs_tokio diff --git a/src/rust/cxx/kj-rs/BUILD.bazel b/src/rust/cxx/kj-rs/BUILD.bazel index 971b7e12512..a80c2a575bb 100644 --- a/src/rust/cxx/kj-rs/BUILD.bazel +++ b/src/rust/cxx/kj-rs/BUILD.bazel @@ -11,7 +11,10 @@ wd_cc_library( "@platforms//os:windows": True, "//conditions:default": False, }), - visibility = ["//src/rust/cxx/tests:__pkg__"], + visibility = [ + "//src/rust/cxx/kj-rs-tokio:__pkg__", + "//src/rust/cxx/tests:__pkg__", + ], deps = [ ":bridge", ], @@ -54,13 +57,17 @@ rust_test( rust_cxx_bridge( name = "bridge", - src = "lib.rs", + src = "ffi.rs", hdrs = glob(["*.h"]), include_prefix = "kj-rs", visibility = ["//src/rust/cxx/tests:__pkg__"], deps = [ "//src/rust/cxx:core", "@capnp-cpp//src/kj:kj", - "@capnp-cpp//src/kj:kj-async", + # kj-rs is the base cxx<->rust Promise/Future bridge: it uses only the abstract async + # core (kj::Promise / kj::EventLoop via async.h), no kj OS I/O. Depending on + # :kj-async-core (not the :kj-async umbrella) keeps the whole kj-rs stack -- and thus + # kj-rs-io / kj-rs-tokio built on it -- off the concrete kj OS event loop (:kj-async-os). + "@capnp-cpp//src/kj:kj-async-core", ], ) diff --git a/src/rust/cxx/kj-rs/awaiter.c++ b/src/rust/cxx/kj-rs/awaiter.c++ index 7bd3aa0fafb..d2e62ee2176 100644 --- a/src/rust/cxx/kj-rs/awaiter.c++ +++ b/src/rust/cxx/kj-rs/awaiter.c++ @@ -1,6 +1,6 @@ #include "awaiter.h" -#include +#include #include @@ -28,13 +28,29 @@ RustPromiseAwaiter::RustPromiseAwaiter( } RustPromiseAwaiter::~RustPromiseAwaiter() noexcept(false) { - // Our `tracePromise()` implementation checks for a null `node`, so we don't have to sever our - // LinkedGroup relationship before destroying `node`. If our FuturePollEvent (our LinkedGroup) - // tries to trace us between now and our destructor completing, `tracePromise()` will ignore the - // null `node`. + // Sever our weak link to any FuturePollEvent before we go away, so it can't trace into or arm a + // destroyed awaiter. Our `tracePromise()` implementation also checks for a null `node`, so even + // between clearPollEvent() and node reset we are safe to trace. + clearPollEvent(); unwindDetector.catchExceptionsIfUnwinding([this]() { node = nullptr; }); } +void RustPromiseAwaiter::setPollEvent(FuturePollEvent& futurePollEvent) { + KJ_IF_SOME(old, maybePollEvent) { + if (&old == &futurePollEvent) return; + old.leaves.remove(*this); + } + futurePollEvent.leaves.add(*this); + maybePollEvent = futurePollEvent; +} + +void RustPromiseAwaiter::clearPollEvent() { + KJ_IF_SOME(old, maybePollEvent) { + old.leaves.remove(*this); + maybePollEvent = kj::none; + } +} + void RustPromiseAwaiter::fire() { // Safety: Our Event can only fire on the event loop which was active when our Event base class // was constructed. Therefore, we don't need to check that we're on the correct event loop. @@ -42,10 +58,10 @@ void RustPromiseAwaiter::fire() { // Nullify our `maybeOptionWaker` to signal that we are done. KJ_DEFER(maybeOptionWaker = kj::none); - KJ_IF_SOME(futurePollEvent, linkedGroup().tryGet()) { + KJ_IF_SOME(futurePollEvent, maybePollEvent) { // Optimized path: we're still linked to a FuturePollEvent. Arm it directly. futurePollEvent.armDepthFirst(); - linkedGroup().set(kj::none); + clearPollEvent(); } else KJ_IF_SOME(optionWaker, maybeOptionWaker) { // We use wake_if_some() rather than an unconditional wake because the OptionWaker may be empty. This // happens when poll() took the optimized path (clearing the OptionWaker and linking to a @@ -68,7 +84,7 @@ void RustPromiseAwaiter::traceEvent(kj::_::TraceBuilder& builder) { node->tracePromise(builder, true); } // TODO(someday): Can we add an entry for the `.await` expression in Rust here? - KJ_IF_SOME(futurePollEvent, linkedGroup().tryGet()) { + KJ_IF_SOME(futurePollEvent, maybePollEvent) { futurePollEvent.traceEvent(builder); } } @@ -82,48 +98,21 @@ void RustPromiseAwaiter::tracePromise(kj::_::TraceBuilder& builder, bool stopAtN // TODO(someday): Can we add an entry for the `.await` expression in Rust here? } -bool RustPromiseAwaiter::poll(const WakerRef& waker, const KjWaker* maybeKjWaker) { +bool RustPromiseAwaiter::poll(const WakerRef& waker) { // TODO(perf): If `this->isNext()` is true, meaning our event is next in line to fire, can we // disarm it, set `done = true`, etc.? If we can only suspend if our enclosing KJ coroutine has - // suspended at least once, we may be able to check for that through LazyArcWaker, but this path + // suspended at least once, we may be able to check for that through PollWaker, but this path // doesn't have access to one. KJ_IF_SOME(optionWaker, maybeOptionWaker) { // Our Promise is not yet ready. - // Check for an optimized wake path. - KJ_IF_SOME(kjWaker, maybeKjWaker) { - KJ_IF_SOME(futurePollEvent, kjWaker.tryGetFuturePollEvent()) { - // Optimized path. The Future which is polling our Promise is in turn being polled by a - // `co_await` expression somewhere up the stack from us. We can arrange to arm the - // `co_await` expression's KJ Event directly when our Promise is ready. - - // Drop any Waker stored in OptionWaker. We'll use the LinkedGroup to wake instead. - // - // Note: this leaves OptionWaker empty while maybeOptionWaker is still Some(ref). If the - // FuturePollEvent is later destroyed (severing the LinkedGroup link) before our Promise - // fires, fire() will find no LinkedGroup AND an empty OptionWaker. fire() handles this - // via wake_if_some(), which is a no-op on an empty OptionWaker. - optionWaker.set_none(); - - // Store a reference to the current `co_await` expression's Future polling Event. The - // reference is weak, and will be cleared if the `co_await` expression happens to end before - // our Promise is ready. In the more likely case that our Promise becomes ready while the - // `co_await` expression is still active, we'll arm its Event so it can `poll()` us again. - linkedGroup().set(futurePollEvent); - - return false; - } - } - - // Unoptimized fallback path. - // Tell our OptionWaker to store a clone of whatever Waker we were given. optionWaker.set(waker); - // Clearing our reference to the FuturePollEvent (if we have one) tells our fire() + // Clearing our weak reference to the FuturePollEvent (if we have one) tells our fire() // implementation to use our OptionWaker to perform the wake. - linkedGroup().set(kj::none); + clearPollEvent(); return false; } else { @@ -132,6 +121,40 @@ bool RustPromiseAwaiter::poll(const WakerRef& waker, const KjWaker* maybeKjWaker } } +bool RustPromiseAwaiter::poll(const WakerRef& waker, const PollWaker& pollWaker) { + KJ_IF_SOME(futurePollEvent, pollWaker.tryGetFuturePollEvent()) { + KJ_IF_SOME(optionWaker, maybeOptionWaker) { + // Our Promise is not yet ready, and we have an optimized wake path. The Future which is + // polling our Promise is in turn being polled by a `co_await` expression somewhere up the + // stack from us. We can arrange to arm the `co_await` expression's KJ Event directly when + // our Promise is ready. + + // Drop any Waker stored in OptionWaker. We'll use our weak link to the FuturePollEvent to + // wake instead. + // + // Note: this leaves OptionWaker empty while maybeOptionWaker is still Some(ref). If the + // FuturePollEvent is later destroyed (severing our weak link) before our Promise fires, + // fire() will find no linked FuturePollEvent AND an empty OptionWaker. fire() handles this + // via wake_if_some(), which is a no-op on an empty OptionWaker. + optionWaker.set_none(); + + // Store a weak reference to the current `co_await` expression's Future polling Event. The + // reference is weak, and will be cleared if the `co_await` expression happens to end before + // our Promise is ready. In the more likely case that our Promise becomes ready while the + // `co_await` expression is still active, we'll arm its Event so it can `poll()` us again. + setPollEvent(futurePollEvent); + + return false; + } else { + // Our Promise is ready. + return true; + } + } + // The PollWaker exposes no FuturePollEvent (its owning thread's kj::Executor is not ours -- + // cannot normally happen in the single-thread world). Fall back to the generic path. + return poll(waker); +} + OwnPromiseNode RustPromiseAwaiter::take_own_promise_node() { KJ_ASSERT(maybeOptionWaker == kj::none, "take_own_promise_node() should only be called after poll() " @@ -151,43 +174,29 @@ void guarded_rust_promise_awaiter_drop_in_place(GuardedRustPromiseAwaiter* ptr) // ======================================================================================= // FuturePollEvent -void FuturePollEvent::exitPollScope(kj::Maybe> maybePromise) { - // Await any LazyArcWaker promise that got created during the call to `poll()`. Note that if a - // Future returns Ready _and_ synchronously wakes its Waker, the work done to await the - // LazyArcWaker promise is wasted, since we will immediately tear the entire BoxFutureAwaiter - // down. However, that's an unlikely case, and this work here isn't likely to be a significant - // source of overhead. - KJ_IF_SOME(promise, maybePromise) { - auto& node = arcWakerPromise.emplace(kj::_::PromiseNode::from(kj::mv(promise))); - node->setSelfPointer(&node); - node->onReady(this); +FuturePollEvent::~FuturePollEvent() noexcept(false) { + // Our FutureWakerCell (if any) is neutralized by the wakerCell guard's destructor during member + // destruction, so any waker reference Rust retained past our lifetime observes a dead weak link + // on a later wake and is a safe no-op, rather than arming this freed event. + + // Sever our weak links to all leaves, so a RustPromiseAwaiter that outlives us (e.g. a stashed + // PromiseFuture) never arms this freed event. + for (;;) { + auto it = leaves.begin(); + if (it == leaves.end()) break; + auto& leaf = *it; + leaves.remove(leaf); + leaf.maybePollEvent = kj::none; } } -void FuturePollEvent::enterPollScope() noexcept { - // Clear out any previous LazyArcWaker promise the FuturePollEvent was holding onto. Note that - // since there is no code path which rejects this Promise, this is not strictly required for - // correctness, but nevertheless serves as a useful assertion. - KJ_IF_SOME(node, arcWakerPromise) { - kj::_::ExceptionOr output; - - node->get(output); - KJ_IF_SOME(exception, kj::runCatchingExceptions([this]() { arcWakerPromise = kj::none; })) { - output.addException(kj::mv(exception)); - } - - // NOTE: `node` is now dangling. - - KJ_IF_SOME(exception, output.exception) { - // We should only ever receive a WakeInstruction, never an exception. If we do receive an - // exception, it would be because our ArcWaker implementation allowed its cross-thread promise - // fulfiller to be destroyed without being fulfilled, or because we foolishly added an - // explicit call to the fulfiller's reject() function. Either way, it is a programming error, - // so we abort the process here by re-throwing across a noexcept boundary. This avoids having - // implement the ability to "reject" the Future poll() Event. - kj::throwFatalException(kj::mv(exception)); - } +kj::Rc FuturePollEvent::cloneWakerCell() { + // Lazily create the cell, bound to this event. + if (wakerCell.cell == nullptr) { + wakerCell.cell = kj::rc(*this); } + // Hand out a new strong reference for Rust to retain. + return wakerCell.cell.addRef(); } void FuturePollEvent::tracePromise(kj::_::TraceBuilder& builder, bool stopAtNextEvent) { @@ -199,32 +208,9 @@ void FuturePollEvent::tracePromise(kj::_::TraceBuilder& builder, bool stopAtNext // When tracing, we can only pick one branch to follow. Arbitrarily, I'm following the first // RustPromiseAwaiter branch, similar to how ExclusiveJoinPromiseNode chooses its left branch. In // the common case, this will be whatever OwnPromiseNode our Rust Future is currently `.await`ing. - auto rustPromiseAwaiters = linkedObjects(); - if (rustPromiseAwaiters.begin() != rustPromiseAwaiters.end()) { + if (!leaves.empty()) { // Our Rust Future is awaiting an OwnPromiseNode. We'll pick the first one in our list. - rustPromiseAwaiters.front().tracePromise(builder, false); - } else KJ_IF_SOME(node, arcWakerPromise) { - // Our Rust Future is not awaiting any OwnPromiseNode, and instead cloned our Waker. We'll trace - // our ArcWaker Promise instead. - if (node.get() != nullptr) { - node->tracePromise(builder, false); - } - } -} - -FuturePollEvent::PollScope::PollScope(FuturePollEvent& futurePollEvent): holder(futurePollEvent) { - futurePollEvent.enterPollScope(); -} - -FuturePollEvent::PollScope::~PollScope() noexcept(false) { - holder.get().futurePollEvent.exitPollScope(reset()); -} - -kj::Maybe FuturePollEvent::PollScope::tryGetFuturePollEvent() const { - KJ_IF_SOME(h, holder.tryGet()) { - return h.futurePollEvent; - } else { - return kj::none; + leaves.front().tracePromise(builder, false); } } diff --git a/src/rust/cxx/kj-rs/awaiter.h b/src/rust/cxx/kj-rs/awaiter.h index ec707c2753a..5f121d3cde9 100644 --- a/src/rust/cxx/kj-rs/awaiter.h +++ b/src/rust/cxx/kj-rs/awaiter.h @@ -1,19 +1,20 @@ #pragma once #include "kj-rs/executor-guarded.h" -#include "kj-rs/linked-group.h" +#include "kj-rs/promise.h" #include "kj-rs/waker.h" #include +#include namespace kj_rs { // ======================================================================================= // Opaque Rust types // -// The following types are defined in lib.rs, and thus in lib.rs.h. lib.rs.h depends on our C++ -// headers, including awaiter.h (the file you're currently reading), so we forward-declare some types -// here for use in the C++ headers. +// The following types are defined in the cxx bridge (ffi.rs), and thus in ffi.rs.h. ffi.rs.h +// depends on our C++ headers, including awaiter.h (the file you're currently reading), so we +// forward-declare some types here for use in the C++ headers. // Wrapper around an `&std::task::Waker`, passed to `RustPromiseAwaiter::poll()`. This indirection // is required because cxx-rs does not permit us to expose opaque Rust types to C++ defined outside @@ -46,17 +47,17 @@ struct OptionWaker; // alignment using bindgen. See inside awaiter.c++ for a static_assert to remind us to re-run // bindgen. // -// RustPromiseAwaiter has two base classes: KJ Event, and a LinkedObject template instantiation. We -// use the Event to discover when our wrapped Promise is ready. Our Event fire() implementation -// records the fact that we are done, then wakes our Waker or arms the FuturePollEvent, if we -// have one. We access the FuturePollEvent via our LinkedObject base class mixin. It gives us the -// ability to store a weak reference to the FuturePollEvent, if we were last polled by one. +// RustPromiseAwaiter has one base class: KJ Event. We use the Event, via the native +// `node->onReady(this)` mechanism, to discover when our wrapped Promise is ready (exactly as a KJ +// coroutine registers its own Event on the promise it `co_await`s). Our Event fire() implementation +// records the fact that we are done, then wakes our Waker or arms the FuturePollEvent, if we have +// one. We hold a weak reference to the FuturePollEvent that last polled us (see maybePollEvent +// below), so a fired Promise can arm that poll event directly. // // Cancellation: Dropping the RustPromiseAwaiter destroys its OwnPromiseNode, cancelling the // wrapped KJ promise. If the RustPromiseAwaiter was never constructed, Rust's OwnPromiseNode::drop() // cancels the promise directly. -class RustPromiseAwaiter final: public kj::_::Event, - public LinkedObject { +class RustPromiseAwaiter final: public kj::_::Event { public: // The Rust code which constructs RustPromiseAwaiter passes us a pointer to a OptionWaker, which can // be thought of as a Rust-native component RustPromiseAwaiter. Its job is to hold a clone of @@ -84,18 +85,29 @@ class RustPromiseAwaiter final: public kj::_::Event, // Poll this Promise for readiness. // - // If the Waker is a KjWaker, you may pass the KjWaker pointer as a second parameter. This may - // allow the implementation of `poll()` to optimize the wake by arming a KJ Event directly when - // the wrapped Promise becomes ready. - // - // If the Waker is not a KjWaker, the `maybeKjWaker` pointer argument must be nullptr. - bool poll(const WakerRef& waker, const KjWaker* maybeKjWaker); + // The two-argument overload is for polls driven by a PollWaker (i.e. a `co_await`ed Future's + // poll): it may optimize the wake by arming a KJ Event directly when the wrapped Promise + // becomes ready. Polls driven by any other Waker use the one-argument overload. + bool poll(const WakerRef& waker); + bool poll(const WakerRef& waker, const PollWaker& pollWaker); // Release ownership of the OwnPromiseNode. Asserts if called before the Promise is ready; that // is, `poll()` must have returned true prior to calling `take_own_promise_node()`. OwnPromiseNode take_own_promise_node(); private: + // Purpose-built one-to-many weak link to the FuturePollEvent that last polled us. The link is + // weak in both directions: destroying either side severs it (see + // clearPollEvent() and ~FuturePollEvent()), so a fired Promise never arms, and tracing never + // touches, a destroyed FuturePollEvent. A RustPromiseAwaiter may outlive the FuturePollEvent that + // first polled it (e.g. a stashed PromiseFuture) and later re-link to a different one. + friend class FuturePollEvent; + void setPollEvent(FuturePollEvent& futurePollEvent); + void clearPollEvent(); + + kj::Maybe maybePollEvent; + kj::ListLink link; + // The Rust code which instantiates RustPromiseAwaiter does so with a OptionWaker object right // next to the RustPromiseAwaiter, such that it is dropped after RustPromiseAwaiter. Thus, our // reference to our OptionWaker is stable. We use the OptionWaker to (optionally) store a clone of @@ -117,8 +129,11 @@ struct GuardedRustPromiseAwaiter: ExecutorGuarded { // We need to inherit constructors or else placement-new will try to aggregate-initialize us. using ExecutorGuarded::ExecutorGuarded; - bool poll(const WakerRef& waker, const KjWaker* maybeKjWaker) { - return get().poll(waker, maybeKjWaker); + bool poll(const WakerRef& waker) { + return get().poll(waker); + } + bool pollWithPollWaker(const WakerRef& waker, const PollWaker& pollWaker) { + return get().poll(waker, pollWaker); } OwnPromiseNode take_own_promise_node() { return get().take_own_promise_node(); @@ -136,20 +151,20 @@ void guarded_rust_promise_awaiter_drop_in_place(GuardedRustPromiseAwaiter*); // `Event::fire()` override which actually polls the Future; this class implements all other base // class virtual functions. // -// A FuturePollEvent contains an optional ArcWakerPromiseAwaiter and a list of zero or more -// RustPromiseAwaiters. These "sub-Promise awaiters" all wrap a KJ Promise of some sort, and arrange -// to arm the FuturePollEvent when their Promises become ready. +// A FuturePollEvent owns an optional FutureWakerCell (handed out to Rust by PollWaker::cloneCell()) +// and a list of zero or more RustPromiseAwaiters. These "sub-Promise awaiters" all wrap a KJ +// Promise of some sort, and arrange to arm the FuturePollEvent when their Promises become ready; a +// woken FutureWakerCell arms it the same way. // // The PromiseNode base class is a hack to implement async tracing. That is, we only implement the // `tracePromise()` function, and decide which Promise to trace into if/when the coroutine calls our // `tracePromise()` implementation. This primarily makes the lifetimes easier to manage: our -// RustPromiseAwaiter LinkedObjects have independent lifetimes from the FuturePollEvent, so we -// mustn't leave references to them, or their members, lying around in the Coroutine class. -class FuturePollEvent: public kj::_::PromiseNode, - public kj::_::Event, - public LinkedGroup { +// weakly-linked RustPromiseAwaiter leaves have independent lifetimes from the FuturePollEvent, so +// we mustn't leave references to them, or their members, lying around in the Coroutine class. +class FuturePollEvent: public kj::_::PromiseNode, public kj::_::Event { public: FuturePollEvent(kj::SourceLocation location = {}): Event(location) {} + ~FuturePollEvent() noexcept(false); // ------------------------------------------------------- // PromiseNode API @@ -159,42 +174,34 @@ class FuturePollEvent: public kj::_::PromiseNode, void tracePromise(kj::_::TraceBuilder& builder, bool stopAtNextEvent) override; - protected: - // PollScope is a LazyArcWaker which is associated with a specific FuturePollEvent, allowing - // optimized Promise `.await`s. Additionally, PollScope's destructor arranges to await any - // ArcWaker promise which was lazily created. - // - // Used by FutureAwaiter, our derived class. - class PollScope; - - private: - // Private API for PollScope. - void enterPollScope() noexcept; - void exitPollScope(kj::Maybe> maybeLazyArcWakerPromise); - - kj::Maybe arcWakerPromise; -}; - -class FuturePollEvent::PollScope: public LazyArcWaker { - public: - // `futurePollEvent` is the FuturePollEvent responsible for calling `Future::poll()`, and must - // outlive this PollScope. - PollScope(FuturePollEvent& futurePollEvent); - ~PollScope() noexcept(false); - KJ_DISALLOW_COPY_AND_MOVE(PollScope); - - // The Event which is using this PollScope to poll() a Future. Waking this FuturePollEvent's - // PollScope arms this Event (possibly via a cross-thread promise fulfiller). We also arm the - // Event directly in the RustPromiseAwaiter class, to more optimally `.await` KJ Promises from - // within Rust. If the current thread's kj::Executor is not the same as the one which owns the - // FuturePollEvent, this function returns kj::none. - kj::Maybe tryGetFuturePollEvent() const override; - private: - struct FuturePollEventHolder { - FuturePollEvent& futurePollEvent; + // Get-or-create this event's FutureWakerCell and hand out a new strong reference to it. Used + // by PollWaker::cloneCell(). The cell is created lazily on first clone and lives until both + // this event and every Rust reference are gone; ~FuturePollEvent neutralizes it so late wakes + // no-op. + friend class PollWaker; + kj::Rc cloneWakerCell(); + + // Weakly-linked list of the RustPromiseAwaiters ("leaves") this Future is currently `.await`ing + // and which may arm this poll event when their Promises become ready. Severed on destruction so + // a leaf that outlives us never arms a freed event. + friend class RustPromiseAwaiter; + kj::List leaves; + + // The FutureWakerCell handed out by cloneWakerCell(), null until the first clone. We hold one + // strong reference through this guard, whose destructor neutralizes the cell (nulling its weak + // Event link so retained Rust references become safe no-ops) before releasing it — the + // invalidation is tied to this event's destruction structurally, not by a destructor body + // remembering to call it. + struct NeutralizeGuard { + kj::Rc cell; + ~NeutralizeGuard() noexcept(false) { + if (cell != nullptr) { + cell->neutralize(); + } + } }; - ExecutorGuarded holder; + NeutralizeGuard wakerCell; }; // ======================================================================================= @@ -204,7 +211,7 @@ template concept Future = requires(F f) { typename F::Output; { - f.poll(kj::instance(), + f.poll(kj::instance(), kj::instance&>()) } -> std::same_as; }; @@ -256,15 +263,10 @@ class FutureAwaiter final: public FuturePollEvent { void poll() { if (isDone()) return; - // TODO(perf): Check if we already have an ArcWaker from a previous suspension and give it to - // LazyArcWaker for cloning if we have the last reference to it at this point. This could save - // memory allocations, but would depend on making XThreadFulfiller and XThreadPaf resettable - // to really benefit. - { - PollScope pollScope(*this); + PollWaker pollWaker(*this); - future.poll(pollScope, result); + future.poll(pollWaker, result); if (isDone()) { onReadyEvent.arm(); } diff --git a/src/rust/cxx/kj-rs/awaiter.rs b/src/rust/cxx/kj-rs/awaiter.rs index 9c32cc6faf3..706ea604099 100644 --- a/src/rust/cxx/kj-rs/awaiter.rs +++ b/src/rust/cxx/kj-rs/awaiter.rs @@ -1,3 +1,8 @@ +//! FFI island (see crate-root `#![deny(unsafe_code)]`): placement bridge for the C++ +//! `GuardedRustPromiseAwaiter` and the `.await` poll glue — Pin projection and placement +//! new/drop over Rust-owned memory. A genuine unsafe seam. +#![allow(unsafe_code)] + use std::mem::MaybeUninit; use std::pin::Pin; use std::task::Context; @@ -7,7 +12,7 @@ use crate::OwnPromiseNode; // Await syntax for OwnPromiseNode use crate::ffi::GuardedRustPromiseAwaiter; use crate::ffi::GuardedRustPromiseAwaiterRepr; -use crate::waker::try_into_kj_waker_ptr; +use crate::waker::try_poll_waker; pub struct PromiseAwaiter { node: Option, @@ -17,6 +22,16 @@ pub struct PromiseAwaiter { // Safety: `option_waker` must be declared after `awaiter`, because `awaiter` contains a reference // to `option_waker`. This ensures `option_waker` will be dropped after `awaiter`. option_waker: OptionWaker, + // Suppresses the auto `Unpin` impl (for this type and any wrapper like `PromiseFuture`). + // After the first poll, `awaiter` holds an in-place-constructed C++ object that is (a) a + // `kj::_::Event` registered with the event loop, (b) the target of the promise node's + // self-pointer (`setSelfPointer` points INTO this memory), (c) the holder of a reference to + // the sibling `option_waker` field, and (d) possibly threaded into a `FuturePollEvent`'s + // intrusive `leaves` list. Moving `self` after that (e.g. `let g = f;` after a `&mut f` + // partial await, which `Unpin` would permit in safe code) leaves all four pointers dangling + // -- use-after-free when the promise fires. `PhantomPinned` turns that into a compile error; + // ordinary `.await` pins structurally and is unaffected. + _pinned: std::marker::PhantomPinned, } impl PromiseAwaiter { @@ -27,6 +42,7 @@ impl PromiseAwaiter { awaiter: MaybeUninit::uninit(), awaiter_initialized: false, option_waker: OptionWaker::empty(), + _pinned: std::marker::PhantomPinned, } } @@ -45,6 +61,12 @@ impl PromiseAwaiter { // contents into GuardedRustPromiseAwaiter's constructor. On all subsequent invocations, `node` // will be None and the constructor will not run. let node = this.node.take(); + // `node` is `Some` on this first (initializing) invocation; see the comment above. + #[expect( + clippy::expect_used, + reason = "get_awaiter initializes exactly once while node is Some (awaiter_initialized guards re-entry); None here is an unreachable internal-invariant violation" + )] + let node = node.expect("node should be Some in call to init()"); // Safety: `awaiter` stores `rust_waker_ptr` and uses it to call `wake()`. Note that // `awaiter` is `this.awaiter`, which lives before `this.option_waker`. @@ -64,7 +86,7 @@ impl PromiseAwaiter { .as_mut_ptr() .cast::(), rust_waker_ptr, - node.expect("node should be Some in call to init()"), + node, ); } this.awaiter_initialized = true; @@ -81,20 +103,25 @@ impl PromiseAwaiter { } pub fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> bool { - let maybe_kj_waker = try_into_kj_waker_ptr(cx.waker()); - let awaiter = self.as_mut().get_awaiter(); - // Safety: The awaiter is initialized by `get_awaiter()` above. `WakerRef` borrows the - // context's waker, which is alive for the duration of the call. `maybe_kj_waker` is null - // or points to the KjWaker inside the waker (validated by `try_into_kj_waker_ptr`). - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe { awaiter.poll(&WakerRef(cx.waker()), maybe_kj_waker) } + // If the Waker driving this poll lends out a C++ PollWaker, take the optimized entry + // point, which may arm the enclosing `co_await`'s KJ Event directly. Both borrows live + // off `cx` for the duration of the call. + match try_poll_waker(cx.waker()) { + Some(poll_waker) => self + .as_mut() + .get_awaiter() + .poll_with_poll_waker(&WakerRef(cx.waker()), poll_waker), + None => self.as_mut().get_awaiter().poll(&WakerRef(cx.waker())), + } } } impl Drop for PromiseAwaiter { fn drop(&mut self) { if self.awaiter_initialized { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `awaiter_initialized` is true, so `self.awaiter` holds a + // `GuardedRustPromiseAwaiter` constructed in place by `get_awaiter`; drop it in + // place exactly once here (this is the only drop path, and `self` is being dropped). unsafe { crate::ffi::guarded_rust_promise_awaiter_drop_in_place( self.awaiter diff --git a/src/rust/cxx/kj-rs/executor-guarded.c++ b/src/rust/cxx/kj-rs/executor-guarded.c++ index a026dfef756..71b5cc93494 100644 --- a/src/rust/cxx/kj-rs/executor-guarded.c++ +++ b/src/rust/cxx/kj-rs/executor-guarded.c++ @@ -4,12 +4,35 @@ namespace kj_rs { +namespace { + +const kj::Executor* tryGetCurrentThreadExecutor() { + // kj/async.h exposes no "maybe" variant of getCurrentThreadExecutor() (verified: only the + // KJ_REQUIRE-ing accessor exists, and the thread-local EventLoop pointer it checks is private + // to async.c++), so probe by catching its recoverable "No event loop is running on this + // thread" requirement failure and treating it as "no executor". This runs on teardown paths + // only (see requireCurrentOrTearingDown below), so the exception cost is irrelevant. + const kj::Executor* current = nullptr; + auto maybeException = + kj::runCatchingExceptions([&]() { current = &kj::getCurrentThreadExecutor(); }); + // (maybeException != kj::none) <=> no live EventLoop on this thread; leave `current` null. + (void)maybeException; + return current; +} + +} // namespace + bool isCurrent(const kj::Executor& executor) { - return &executor == &kj::getCurrentThreadExecutor(); + return tryGetCurrentThreadExecutor() == &executor; } void requireCurrent(const kj::Executor& executor, kj::LiteralStringConst message) { KJ_REQUIRE(isCurrent(executor), message); } +void requireCurrentOrTearingDown(const kj::Executor& executor, kj::LiteralStringConst message) { + const kj::Executor* current = tryGetCurrentThreadExecutor(); + KJ_REQUIRE(current == &executor || current == nullptr, message); +} + } // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/executor-guarded.h b/src/rust/cxx/kj-rs/executor-guarded.h index e7f68c8d7d9..b00fb2944ca 100644 --- a/src/rust/cxx/kj-rs/executor-guarded.h +++ b/src/rust/cxx/kj-rs/executor-guarded.h @@ -4,11 +4,20 @@ namespace kj_rs { -// Return true if `executor`'s event loop is active on the current thread. +// Return true if `executor`'s event loop is active on the current thread. Never throws: a +// thread with no live event loop at all (e.g. after ~EventLoop during teardown) reports false. bool isCurrent(const kj::Executor& executor); // Assert that `executor`'s event loop is active on the current thread, or throw an exception // containing `message`. void requireCurrent(const kj::Executor& executor, kj::LiteralStringConst message); +// Like requireCurrent(), but tolerant of full loop teardown: if the current thread has NO event +// loop at all (~EventLoop already ran), return quietly instead of throwing. Used by +// ~ExecutorGuarded, which can legitimately run after loop teardown from Rust drop glue (e.g. the +// tokio runtime cancelling still-pending LocalSet tasks in TokioPort::drop, after +// TokioAsyncIoContext destroyed the WaitScope and EventLoop); a throw there would unwind through +// a cxx `prevent_unwind` boundary and abort the process. Destruction on a thread running a +// DIFFERENT live event loop is still a contract violation and throws. +void requireCurrentOrTearingDown(const kj::Executor& executor, kj::LiteralStringConst message); // ExecutorGuarded is a helper class which allows mutable access to a wrapped value to any thread // running the KJ event loop that was active at the time of construction. Any access attempts by a @@ -19,7 +28,10 @@ class ExecutorGuarded { template ExecutorGuarded(Args&&... args): value(kj::fwd(args)...) {} ~ExecutorGuarded() noexcept(false) { - requireCurrent(executor, "destruction on wrong event loop"_kjc); + // Teardown-tolerant (see requireCurrentOrTearingDown above): destruction with no event loop + // on the thread proceeds quietly (best-effort destruction of `value`); destruction on a + // thread running a different live loop still throws. + requireCurrentOrTearingDown(*executor, "destruction on wrong event loop"_kjc); } KJ_DISALLOW_COPY_AND_MOVE(ExecutorGuarded); @@ -29,7 +41,7 @@ class ExecutorGuarded { // Throws an exception with `message` if the current thread is not running the expected event // loop. T& get(kj::LiteralStringConst message = "access on wrong event loop"_kjc) const { - requireCurrent(executor, message); + requireCurrent(*executor, message); // Safety: const_cast is okay because we know that we are being accessed on a thread running our // original event loop. All successful accesses through `get()` are effectively single-threaded, @@ -38,7 +50,7 @@ class ExecutorGuarded { } kj::Maybe tryGet() const { - if (isCurrent(executor)) { + if (isCurrent(*executor)) { // Safety: const_cast is okay because we know that we are being accessed on a thread running our // original event loop. All successful accesses through `get()` are effectively single-threaded, // even though the event loop, and this object, may collectively move between threads. @@ -49,7 +61,15 @@ class ExecutorGuarded { } private: - const kj::Executor& executor = kj::getCurrentThreadExecutor(); + // Owned (via `addRef()`) rather than a bare `const kj::Executor&`, so the Executor object stays + // alive as long as this guard does. A guarded value can be destroyed by Rust *after* the KJ + // event loop has been torn down (e.g. a bridged future dropped during teardown); with a bare + // reference the destructor's executor check would take the address of — and a reused address + // could alias — a freed Executor. `addRef()` keeps the Executor at a stable, valid address; if + // the loop is gone, `isCurrent` reports false without throwing and the destructor's + // `requireCurrentOrTearingDown` lets destruction proceed quietly (`get()`/`requireCurrent` + // still throw, since post-teardown *access* remains a contract violation). + kj::Own executor = kj::getCurrentThreadExecutor().addRef(); T value; }; diff --git a/src/rust/cxx/kj-rs/ffi.rs b/src/rust/cxx/kj-rs/ffi.rs new file mode 100644 index 00000000000..51917b8b7f9 --- /dev/null +++ b/src/rust/cxx/kj-rs/ffi.rs @@ -0,0 +1,135 @@ +//! The `#[cxx::bridge]` FFI island for kj-rs. +//! +//! This is the crate's dedicated `#[cxx::bridge]` file (file-top `#![allow(unsafe_code)]`): the +//! cxx bridge DSL and its generated glue are inherently unsafe (extern "C++" vtables, placement +//! new/drop, C-ABI signatures) — a genuine seam. The bridge module is re-exported as +//! `crate::ffi::*`, the path the rest of the crate (awaiter.rs, promise.rs, waker.rs) uses. +//! +//! The crate root `lib.rs` is wholly-safe. The per-type vocabulary islands +//! (`future.rs`, `awaiter.rs`, `waker.rs`, `promise.rs`, `own.rs`, `refcount.rs`, `maybe.rs`) +//! carry their own file-top `#![allow(unsafe_code)]`: they implement individual unsafe primitive +//! types, distinct from this — the crate's cxx bridge. +#![allow(unsafe_code)] + +pub use bridge::*; + +use crate::awaiter::OptionWaker; +use crate::awaiter::WakerRef; + +#[cxx::bridge(namespace = "kj_rs")] +// The cxx bridge DSL and its generated glue are inherently unsafe (extern "C++" vtables, placement +// new/drop, C-ABI signatures). This is a genuine seam. +// missing_safety_doc: the `# Safety` docs on `reown` below are for human readers only — the cxx +// macro does not forward doc comments to the generated unsafe shim, so the lint cannot be +// satisfied by documentation here. +#[expect(clippy::missing_safety_doc)] +mod bridge { + + /// Representation of a `GuardedRustPromiseAwaiter` in C++. The size of the blob should match. + #[derive(Debug)] + pub struct GuardedRustPromiseAwaiterRepr { + _bindgen_opaque_blob: [u64; 14usize], + } + + extern "Rust" { + type WakerRef<'a>; + } + + extern "Rust" { + // We expose the Rust Waker type to C++ through this OptionWaker reference wrapper. cxx-rs + // does not allow us to export types defined outside this crate, such as Waker, directly. + // + // `LazyRustPromiseAwaiter` (the implementation of `.await` syntax/the IntoFuture trait), + // stores a OptionWaker immediately after `GuardedRustPromiseAwaiter` in declaration order. + // pass the Waker to the `RustPromiseAwaiter` class, which is implemented in C++ + type OptionWaker; + fn set(&mut self, waker: &WakerRef); + fn set_none(&mut self); + fn wake_if_some(&mut self); + } + + unsafe extern "C++" { + include!("kj-rs/waker.h"); + + /// The stack-owned waker C++ passes to `Future::poll()`. Rust only ever borrows it; the + /// `Waker` built from it (waker.rs) has a no-op drop and clones by taking a real strong + /// reference to the event's `FutureWakerCell` via `clone_cell()`. + type PollWaker; + #[cxx_name = "wakeByRef"] + fn wake_by_ref(self: &PollWaker); + #[cxx_name = "cloneCell"] + fn clone_cell(self: &PollWaker) -> KjMaybe>; + + /// The refcounted cell behind every retained waker; waking it arms the owning + /// FuturePollEvent (a safe no-op after that event is destroyed). + type FutureWakerCell; + #[cxx_name = "wakeByRef"] + fn wake_by_ref(self: &FutureWakerCell); + #[cxx_name = "addRef"] + fn add_ref(self: &FutureWakerCell) -> KjRc; + /// Re-own a strong reference previously disowned into a `RawWaker` data slot (waker.rs's + /// owned-cell vtable). + /// + /// # Safety + /// + /// `self` must carry exactly such a surrendered reference — this mints an owned handle + /// without incrementing the count. Dropping the returned handle releases the reference. + unsafe fn reown(self: &FutureWakerCell) -> KjRc; + } + + unsafe extern "C++" { + include!("kj-rs/promise.h"); + + type OwnPromiseNode = crate::OwnPromiseNode; + + // Takes `&mut` (not a raw pointer): this is a placement-destruct of a live + // `OwnPromiseNode` whose backing memory is owned by Rust and only reached through the + // `&mut self` in `OwnPromiseNode`'s `Drop`. The reference is valid for the call; the + // value is logically dead only after, inside `drop`, so no use-after-free is possible. + // Expressing it as a borrow lets cxx generate a safe-to-call binding. + fn own_promise_node_drop_in_place(node: &mut OwnPromiseNode); + } + + unsafe extern "C++" { + include!("kj-rs/awaiter.h"); + + type GuardedRustPromiseAwaiter; + + /// Placement-new of the C++ awaiter into Rust-owned storage. + /// + /// # Safety + /// + /// - `ptr` must point to uninitialized storage of (at least) the size and alignment of + /// `GuardedRustPromiseAwaiterRepr`, valid for writes, and must stay pinned for the + /// awaiter's lifetime. + /// - `rust_waker_ptr` must point to a valid `OptionWaker` that outlives the awaiter (the + /// C++ side stores and dereferences this pointer on later `poll()`s). + /// - The awaiter must eventually be destroyed exactly once via + /// `guarded_rust_promise_awaiter_drop_in_place`. + unsafe fn guarded_rust_promise_awaiter_new_in_place( + ptr: *mut GuardedRustPromiseAwaiter, + rust_waker_ptr: *mut OptionWaker, + node: OwnPromiseNode, + ); + /// Placement-destruct of the awaiter constructed by + /// `guarded_rust_promise_awaiter_new_in_place`. + /// + /// # Safety + /// + /// `ptr` must point to a live awaiter previously constructed in that storage by + /// `guarded_rust_promise_awaiter_new_in_place`, and the awaiter must not be used again + /// afterwards (at most one drop per construction). + unsafe fn guarded_rust_promise_awaiter_drop_in_place(ptr: *mut GuardedRustPromiseAwaiter); + + fn poll(self: Pin<&mut GuardedRustPromiseAwaiter>, waker: &WakerRef) -> bool; + #[cxx_name = "pollWithPollWaker"] + fn poll_with_poll_waker( + self: Pin<&mut GuardedRustPromiseAwaiter>, + waker: &WakerRef, + poll_waker: &PollWaker, + ) -> bool; + + #[must_use] + fn take_own_promise_node(self: Pin<&mut GuardedRustPromiseAwaiter>) -> OwnPromiseNode; + } +} diff --git a/src/rust/cxx/kj-rs/future.h b/src/rust/cxx/kj-rs/future.h index 39aae452a09..78b2abbb5b5 100644 --- a/src/rust/cxx/kj-rs/future.h +++ b/src/rust/cxx/kj-rs/future.h @@ -67,7 +67,7 @@ namespace repr { // ::kj_rs::repr::PollCallback using PollCallback = kj_rs::FuturePollStatus (*)( - void /* RustFuture::fut */* fut, const void* waker, void /* T */* ret); + void /* RustFuture::fut */* fut, const ::kj_rs::PollWaker& waker, void /* T */* ret); // ::kj_rs::repr::DropCallback using DropCallback = void (*)(void /* RustFuture::fut */* fut); @@ -79,8 +79,31 @@ using DropCallback = void (*)(void /* RustFuture::fut */* fut); // which drops the Rust Future and transitively cancels any KJ sub-promises it was .await'ing. struct RustFuture { + // Eager-by-default conversion: the returned promise starts running immediately, without + // being awaited — the future is polled synchronously up to its first suspension point + // (exactly like calling a KJ coroutine, which runs to its first co_await), and continues + // on the event loop from there. KJ code universally assumes promises are "hot" (a stored + // promise still makes progress), so this is the right default for every bridged + // `async fn`; before this conversion was eager, every consumer had to remember a manual + // `.eagerlyEvaluate(nullptr)`. + // + // Requires a current kj::EventLoop on this thread (same requirement as awaiting the + // promise, just enforced at creation). Cancellation is unchanged: dropping the promise + // still synchronously cancels the Rust future and everything it is awaiting. + // + // The rare consumer that genuinely wants a cold future can call `lazily()` below on the + // raw RustFuture instead of going through this conversion (the bridge's generated shims + // always convert eagerly, so that consumer must obtain the RustFuture itself). template operator kj::Promise() { + return lazily().eagerlyEvaluate(nullptr); + } + + // Lazy (cold) conversion: nothing runs until the returned promise is first awaited. + // This is the raw adapter the eager conversion above builds on, and the C++-side escape + // hatch for code that genuinely needs a cold promise. + template + kj::Promise lazily() { struct Impl { using ExceptionOrValue = ::kj::_::ExceptionOr<::kj::_::FixVoid>; using Output = ::kj::_::FixVoid; @@ -100,10 +123,9 @@ struct RustFuture { KJ_DISALLOW_COPY(Impl); - void poll(const ::kj_rs::KjWaker& waker, ExceptionOrValue& output) noexcept { + void poll(const ::kj_rs::PollWaker& waker, ExceptionOrValue& output) noexcept { ::kj_rs::FuturePoller poller; - poller.poll( - [this, &waker](void* result) { return fut.poll(&fut, &waker, result); }, output); + poller.poll([this, &waker](void* result) { return fut.poll(&fut, waker, result); }, output); } RustFuture fut; diff --git a/src/rust/cxx/kj-rs/future.rs b/src/rust/cxx/kj-rs/future.rs index aab7b9dbd32..707eb53618f 100644 --- a/src/rust/cxx/kj-rs/future.rs +++ b/src/rust/cxx/kj-rs/future.rs @@ -1,3 +1,8 @@ +//! FFI island (see crate-root `#![deny(unsafe_code)]`): the `RustFuture` C-ABI vtable that drives +//! all bridged async — `unsafe extern "C"` poll/drop callbacks, raw-pointer result writes, and +//! `Pin`/`Box` raw conversions. A genuine unsafe seam. +#![allow(unsafe_code)] + // This file contains boilerplate which must occur once per crate, rather than once per type. use std::pin::Pin; @@ -33,11 +38,59 @@ pub mod repr { use static_assertions::assert_eq_size; use super::FuturePollStatus; - use crate::KjWaker; + use crate::ffi::PollWaker; + + /// Converts a panic payload (from `std::panic::catch_unwind`) escaping a bridged future + /// into a heap-allocated `kj::Exception` written to the poll callback's output parameter, + /// so C++ observes a rejected promise instead of a process abort. + /// + /// This mirrors the sync bridge path (`cxx::private::try_unwind`/`catch_unwind` in + /// src/unwind.rs), which converts panics in `extern "Rust"` functions into + /// `kj::Exception`s. One divergence: a `cxx::CanceledException` payload (produced when an + /// infallible `extern "C++"` call throws `kj::CanceledException`) cannot be propagated as + /// a distinct "canceled" state here, because `FuturePollStatus` has no Canceled arm; it is + /// reported as a regular `kj::Exception` describing the cancellation instead. + /// + /// # Safety + /// + /// `ret` must point to storage valid for holding a `kj::Exception*` (the C++ + /// `FuturePoller` union guarantees this for the Error arm). + #[expect( + clippy::needless_pass_by_value, + reason = "takes ownership of the panic payload, mirroring std::panic::catch_unwind's Err arm" + )] + unsafe fn write_panic_as_exception( + ret: *mut c_void, + err: Box, + ) -> FuturePollStatus { + let msg = if let Some(s) = err.downcast_ref::<&'static str>() { + format!("panic in bridged future poll: {s}") + } else if let Some(s) = err.downcast_ref::() { + format!("panic in bridged future poll: {s}") + } else if err.downcast_ref::().is_some() { + "panic in bridged future poll: kj::CanceledException".to_owned() + } else { + "panic in bridged future poll".to_owned() + }; + let exception = cxx::IntoKjException::into_kj_exception( + cxx::KjError::new(cxx::KjExceptionType::Failed, msg), + file!(), + line!(), + ); + // SAFETY: `ret` points to storage valid for a `kj::Exception*` per this fn's + // `# Safety` contract (the C++ `FuturePoller` Error arm). + unsafe { + std::ptr::write( + ret.cast::<*mut c_void>(), + exception.into_raw().as_ptr().cast(), + ); + } + FuturePollStatus::ERROR + } - type PollCallback = unsafe extern "C" fn( + type PollCallback = for<'a> unsafe extern "C" fn( fut: *mut c_void, - waker: *const c_void, + waker: &'a PollWaker, ret: *mut c_void, ) -> FuturePollStatus; @@ -71,27 +124,41 @@ pub mod repr { assert_eq_size!(RustInfallibleFuture<()>, [*mut c_void; 4]); impl RustFuture<'_, T> { + /// # Safety + /// + /// C++ `RustFuture` vtable protocol (future.h): `fut` must be the `fut` field of a + /// live, not-yet-dropped `RustFuture` created by [`future`]; `ret` must point to + /// storage suitable for a `T` + /// (Complete) or a `kj::Exception*` (Error), per `FuturePoller`'s union. + /// + /// Unwind safety: any panic escaping the wrapped future's `poll` (or the waker + /// machinery) is caught here and converted into an errored completion, because + /// unwinding out of an `extern "C"` fn is instant process abort (Rust >= 1.81). + /// This makes a panicking bridged `async fn` surface to C++ as a rejected + /// `kj::Promise` carrying a `kj::Exception`, matching the sync bridge path. pub(crate) unsafe extern "C" fn poll( fut: *mut c_void, - waker: *const c_void, + waker: &PollWaker, ret: *mut c_void, ) -> FuturePollStatus { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: per this fn's `# Safety` contract, `fut` is the `fut` field of a live + // `RustFuture`, i.e. a valid `*mut FuturePtr` we may read and then pin. let fut = unsafe { *(fut.cast::>()) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the boxed future is never moved out of its heap allocation, so pinning + // the `&mut` reborrow is sound. let fut = unsafe { Pin::new_unchecked(&mut *fut) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - let waker = unsafe { &*waker.cast::() }; - let waker = Waker::from(waker); - let mut context = Context::from_waker(&waker); - match fut.poll(&mut context) { - Poll::Ready(Ok(value)) => { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let waker = Waker::from(waker); + let mut context = Context::from_waker(&waker); + fut.poll(&mut context) + })) { + Ok(Poll::Ready(Ok(value))) => { + // SAFETY: `ret` points to storage suitable for a `T` (Complete arm). unsafe { std::ptr::write(ret.cast::(), value) }; FuturePollStatus::COMPLETE } - Poll::Ready(Err(error)) => { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + Ok(Poll::Ready(Err(error))) => { + // SAFETY: `ret` points to storage for a `kj::Exception*` (Error arm). unsafe { std::ptr::write( ret.cast::<*mut c_void>(), @@ -100,53 +167,90 @@ pub mod repr { }; FuturePollStatus::ERROR } - Poll::Pending => FuturePollStatus::PENDING, + Ok(Poll::Pending) => FuturePollStatus::PENDING, + // SAFETY: `ret` is the Error-arm storage; forwarded to `write_panic_as_exception`. + Err(panic_payload) => unsafe { write_panic_as_exception(ret, panic_payload) }, } } + /// # Safety + /// + /// C++ `RustFuture` vtable protocol (future.h): `fut` must be the `fut` field of a + /// live `RustFuture` created by [`future`], and must not be used again afterwards + /// (drop-exactly-once, enforced by `Impl`'s move semantics on the C++ side). + /// + /// Unwind safety: a panic in the future's destructor has no error channel (this is + /// called from C++ destructors/cancellation paths), so it is converted into a + /// deterministic, labeled abort via `cxx::private::prevent_unwind` — the same + /// semantics the sync bridge uses for panics that cannot be reported (rather than + /// the unlabeled langdef abort of unwinding out of `extern "C"`). pub(crate) unsafe extern "C" fn drop_in_place(fut: *mut c_void) { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: per this fn's `# Safety` contract, `fut` is the `fut` field of a live + // `RustFuture` not used again, so we may read the pointer, reclaim the box, and pin. let fut = unsafe { *(fut.cast::>()) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `fut` was produced by `Box::into_raw` in [`future`]; reclaim ownership once. let fut = unsafe { Box::from_raw(fut) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the box is never moved out of its allocation, so pinning it is sound. let fut = unsafe { Pin::new_unchecked(fut) }; - drop(fut); + cxx::private::prevent_unwind("kj_rs::repr::RustFuture::drop_in_place", move || { + drop(fut); + }); } } impl RustInfallibleFuture<'_, T> { + /// # Safety + /// + /// Same contract as [`RustFuture::poll`], with `fut` created by [`infallible_future`]. + /// Although the future itself cannot return an error, a panic escaping its `poll` is + /// still converted into an errored completion (`FuturePollStatus::ERROR` writing a + /// `kj::Exception*`): the C++ `FuturePoller` handles the Error arm identically for + /// infallible futures, and `kj::Promise` can always carry an exception. pub(crate) unsafe extern "C" fn poll( fut: *mut c_void, - waker: *const c_void, + waker: &PollWaker, ret: *mut c_void, ) -> FuturePollStatus { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: per this fn's `# Safety` contract, `fut` is the `fut` field of a live + // `RustInfallibleFuture`, i.e. a valid `*mut InfallibleFuturePtr` to read+pin. let fut = unsafe { *(fut.cast::>()) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the boxed future is never moved out of its allocation, so pinning is sound. let fut = unsafe { Pin::new_unchecked(&mut *fut) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - let waker = unsafe { &*waker.cast::() }; - let waker = Waker::from(waker); - let mut context = Context::from_waker(&waker); - match fut.poll(&mut context) { - Poll::Ready(value) => { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let waker = Waker::from(waker); + let mut context = Context::from_waker(&waker); + fut.poll(&mut context) + })) { + Ok(Poll::Ready(value)) => { + // SAFETY: `ret` points to storage suitable for a `T` (Complete arm). unsafe { std::ptr::write(ret.cast::(), value) }; FuturePollStatus::COMPLETE } - Poll::Pending => FuturePollStatus::PENDING, + Ok(Poll::Pending) => FuturePollStatus::PENDING, + // SAFETY: `ret` is the Error-arm storage; forwarded to `write_panic_as_exception`. + Err(panic_payload) => unsafe { write_panic_as_exception(ret, panic_payload) }, } } + /// # Safety + /// + /// Same contract as [`RustFuture::drop_in_place`], with `fut` created by + /// [`infallible_future`]. A panic in the destructor aborts deterministically with a + /// label (see there for rationale). pub(crate) unsafe extern "C" fn drop_in_place(fut: *mut c_void) { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: per this fn's `# Safety` contract, `fut` is the `fut` field of a live + // `RustInfallibleFuture` not used again; read the pointer, reclaim the box, and pin. let fut = unsafe { *(fut.cast::>()) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `fut` came from `Box::into_raw` in [`infallible_future`]; reclaim once. let fut = unsafe { Box::from_raw(fut) }; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the box is never moved out of its allocation, so pinning it is sound. let fut = unsafe { Pin::new_unchecked(fut) }; - drop(fut); + cxx::private::prevent_unwind( + "kj_rs::repr::RustInfallibleFuture::drop_in_place", + move || { + drop(fut); + }, + ); } } @@ -154,7 +258,8 @@ pub mod repr { pub fn future<'a, T: Unpin>( fut: Pin> + 'a>>, ) -> RustFuture<'a, T> { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the box is immediately re-boxed via `Box::into_raw` and only ever reconstituted + // (and re-pinned) in `drop_in_place`, so the pinned future is never moved. let fut = Box::into_raw(unsafe { Pin::into_inner_unchecked(fut) }); let poll = RustFuture::::poll; let drop = RustFuture::::drop_in_place; @@ -165,7 +270,8 @@ pub mod repr { pub fn infallible_future<'a, T: Unpin>( fut: Pin + 'a>>, ) -> RustInfallibleFuture<'a, T> { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the box is immediately re-boxed via `Box::into_raw` and only ever reconstituted + // (and re-pinned) in `drop_in_place`, so the pinned future is never moved. let fut = Box::into_raw(unsafe { Pin::into_inner_unchecked(fut) }); let poll = RustInfallibleFuture::::poll; let drop = RustInfallibleFuture::::drop_in_place; diff --git a/src/rust/cxx/kj-rs/lib.rs b/src/rust/cxx/kj-rs/lib.rs index c7df9fcd7d2..e2a2d995d36 100644 --- a/src/rust/cxx/kj-rs/lib.rs +++ b/src/rust/cxx/kj-rs/lib.rs @@ -1,6 +1,39 @@ -use awaiter::OptionWaker; +// Safety & panic enforcement walls. Inherent-FFI crate: unsafe is +// concentrated at the bridge and every op must sit in an explicit, documented unsafe +// block; prod code returns Result/KjError rather than panicking (a panic on the async +// poll path is a process abort). Test code is exempted below. +#![deny(unsafe_op_in_unsafe_fn)] +// Quarantine unsafe into named FFI islands: deny unsafe crate-wide, then re-allow it only on the +// modules that genuinely need it (each carries its own `#![allow(unsafe_code)]`). Any module +// without that opt-in — and any newly-added module — is compiler-proven unsafe-free, and no future +// edit can smuggle unsafe into non-island code without tripping this deny. +#![deny(unsafe_code)] +#![deny(clippy::undocumented_unsafe_blocks)] +#![deny( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::unreachable, + clippy::todo, + clippy::unimplemented +)] +#![cfg_attr( + test, + allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::unreachable, + clippy::todo, + clippy::unimplemented + ) +)] + +// The cxx bridge expands vocabulary builtins (KjRc, KjMaybe, ...) to `::kj_rs::...` paths; make +// that path resolve inside this crate itself, since ffi.rs's bridge uses them too. +extern crate self as kj_rs; + pub use awaiter::PromiseAwaiter; -use awaiter::WakerRef; pub use date::KjDate; pub use future::FuturePollStatus; pub use future::map_err; @@ -14,10 +47,12 @@ pub use promise::new_callbacks_promise_future; pub use refcount::repr::KjArc; pub use refcount::repr::KjRc; -pub use crate::ffi::KjWaker; +pub use crate::ffi::FutureWakerCell; +pub use crate::ffi::PollWaker; mod awaiter; mod date; +mod ffi; mod future; pub mod maybe; mod own; @@ -36,80 +71,3 @@ pub type Result = std::io::Result; pub type Error = std::io::Error; pub trait JsgStruct {} - -#[cxx::bridge(namespace = "kj_rs")] -mod ffi { - - /// Representation of a `GuardedRustPromiseAwaiter` in C++. The size of the blob should match. - #[derive(Debug)] - pub struct GuardedRustPromiseAwaiterRepr { - _bindgen_opaque_blob: [u64; 13usize], - } - - extern "Rust" { - type WakerRef<'a>; - } - - extern "Rust" { - // We expose the Rust Waker type to C++ through this OptionWaker reference wrapper. cxx-rs - // does not allow us to export types defined outside this crate, such as Waker, directly. - // - // `LazyRustPromiseAwaiter` (the implementation of `.await` syntax/the IntoFuture trait), - // stores a OptionWaker immediately after `GuardedRustPromiseAwaiter` in declaration order. - // pass the Waker to the `RustPromiseAwaiter` class, which is implemented in C++ - type OptionWaker; - fn set(&mut self, waker: &WakerRef); - fn set_none(&mut self); - fn wake_if_some(&mut self); - } - - unsafe extern "C++" { - include!("kj-rs/waker.h"); - - // Match the definition of the abstract virtual class in the C++ header. - type KjWaker; - #[cxx_name = "clone"] - fn clone_kj_waker(&self) -> *const KjWaker; - fn wake(&self); - fn wake_by_ref(&self); - fn drop(&self); - } - - unsafe extern "C++" { - include!("kj-rs/promise.h"); - - type OwnPromiseNode = crate::OwnPromiseNode; - - /// # Safety - /// `node` must point to a live `OwnPromiseNode`. - unsafe fn own_promise_node_drop_in_place(node: *mut OwnPromiseNode); - } - - unsafe extern "C++" { - include!("kj-rs/awaiter.h"); - - type GuardedRustPromiseAwaiter; - - /// # Safety - /// The pointers must identify valid storage and a live waker for the awaiter's lifetime. - unsafe fn guarded_rust_promise_awaiter_new_in_place( - ptr: *mut GuardedRustPromiseAwaiter, - rust_waker_ptr: *mut OptionWaker, - node: OwnPromiseNode, - ); - /// # Safety - /// `ptr` must point to an initialized guarded awaiter. - unsafe fn guarded_rust_promise_awaiter_drop_in_place(ptr: *mut GuardedRustPromiseAwaiter); - - /// # Safety - /// `maybe_kj_waker`, when non-null, must point to a live `KjWaker`. - unsafe fn poll( - self: Pin<&mut GuardedRustPromiseAwaiter>, - waker: &WakerRef, - maybe_kj_waker: *const KjWaker, - ) -> bool; - - #[must_use] - fn take_own_promise_node(self: Pin<&mut GuardedRustPromiseAwaiter>) -> OwnPromiseNode; - } -} diff --git a/src/rust/cxx/kj-rs/linked-group.h b/src/rust/cxx/kj-rs/linked-group.h deleted file mode 100644 index 34567a71e2b..00000000000 --- a/src/rust/cxx/kj-rs/linked-group.h +++ /dev/null @@ -1,315 +0,0 @@ -#pragma once - -#include - -namespace kj_rs { - -// `LinkedGroup` and `LinkedObject` are CRTP mixins which allow derived classes G and O -// to weakly refer to each other in a one-to-many relationship. -// -// For example, say you have two classes, Group and Object. There exists a natural one-to-many -// relationship between the two. Given a Group, you would like to be able to dererefence its -// Objects, and, given an Object, you would like be able to dereference its Group. Further suppose -// the objects have independent lifetimes: Objects may be destroyed before their Groups, and Groups -// may be destroyed before their Objects. -// -// If you are operating in a single-threaded context (or can provide sufficient synchronization), -// and if Group and Object are both immobile (non-copyable, non-moveable) classes, then -// `LinkedGroup` and `LinkedObject` can be used to implement the above -// scenario safely. To do so, first: -// -// - Your Group class must publicly inherit from `LinkedGroup`. -// - Your Object class must publicly inherit from `LinkedObject`. -// -// This will add one protected member function to each of your derived classes: -// `Object::linkedGroup()`, and `Group::linkedObjects()`. They are protected so that they are not -// part of your type's public API unless you explicitly want them to be, e.g., with a public `using` -// statement like `using LinkedGroup::linkedObjects`. -// -// You can use `Object::linkedGroup()` to manage Group membership and dereference Groups from -// Objects: -// -// - `object.linkedGroup().set(group)` adds an Object to a Group. -// This function implicitly removes the Object from its current Group, if any. -// - `object.linkedGroup().set(kj::none)` removes an Object from its current Group, if any. -// - `object.linkedGroup().tryGet()` dereferences the Object's current Group, if any. -// -// You can use `Group::linkedObjects()` to iterate over the list of currently linked Objects. -// -// - `group.linkedObjects().begin()` obtains an iterator to the beginning of the list of Objects. -// - `group.linkedObjects().end()` obtains an iterator to the end of the list of Objets. -// - `group.linkedObjects().front()` dereferences the front of the list of Objects. -// Calling `front()` on an empty list (`begin() == end()`) is undefined behavior. -// - `group.linkedObjects().empty()` is true if there are no Objects in the list. -// -// Finally, destroying either the Group or its Object safely severs their relationship(s). -// -// - Destroying an Object implicitly calls `object.linkedGroup().set(kj::none)` on itself. -// - Destroying a Group implicitly calls `object.linkedGroup().set(kj::none)` on all its objects. -// -// Considerations: -// -// - Your Group object's destructor will contain a _O(n)_ algorithm inside it, with _n_ being the -// number of linked objects at destruction time. If Groups frequently outlive large sets of -// Objects, this may be an issue to consider. -// - It is valid to remove the front Object in a `Group::linkedObjects()` list while iterating -// over the list. Removing an Object in any other position in the list will invalidate all -// existing iterators. -// -// TODO(someday): Multiple inheritance if an object must join multiple groups, or a group must -// have multiple linked object types? Can we write something like `linkedGroup()` in the -// LinkedObject derived class, and `linkedObjects()` in the LinkedGroup derived class? -template -class LinkedGroup; -template -class LinkedObject; - -template -class StaticCastIterator; - -// CRTP mixin for derived class G. -template -class LinkedGroup { - public: - LinkedGroup() = default; - ~LinkedGroup() noexcept(false) { - for (auto& object: list) { - object.removeFromGroup(*this); - } - } - KJ_DISALLOW_COPY_AND_MOVE(LinkedGroup); - - private: - // We'll refer to the `LinkedObject` type quite a bit below, so we shadow the class - // template with our own convenience typedef. But, we need to give LinkedObject friend access to - // us first. - friend class LinkedObject; - using LinkedObject = LinkedObject; - - using List = kj::List; - - using ListIterator = kj::ListIterator; - using ConstListIterator = kj::ListIterator; - - using Iterator = StaticCastIterator; - using ConstIterator = StaticCastIterator; - - protected: - // A proxy class representing this LinkedGroup's list of LinkedObjects, if any. Instead of - // exposing multiple functions on LinkedGroup, we expose one: `linkedObjects()`, and that function - // returns an object of this proxy class (or the similar ConstLinkedObjectList class below). - class LinkedObjectList { - public: - LinkedObjectList(List& list): list(list) {} - Iterator begin() { - return list.begin(); - } - Iterator end() { - return list.end(); - } - decltype(*kj::instance()) front() { - return *begin(); - } - bool empty() const { - return list.empty(); - } - - private: - List& list; - }; - - class ConstLinkedObjectList { - public: - ConstLinkedObjectList(const List& list): list(list) {} - ConstIterator begin() const { - return list.begin(); - } - ConstIterator end() const { - return list.end(); - } - decltype(*kj::instance()) front() const { - return *begin(); - } - bool empty() const { - return list.empty(); - } - - private: - const List& list; - }; - - LinkedObjectList linkedObjects() { - return LinkedObjectList(list); - } - ConstLinkedObjectList linkedObjects() const { - return ConstLinkedObjectList(list); - } - - private: - kj::List list; -}; - -// CRTP mixin for derived class O. -template -class LinkedObject { - public: - LinkedObject() = default; - ~LinkedObject() noexcept(false) { - invalidateGroup(); - } - KJ_DISALLOW_COPY_AND_MOVE(LinkedObject); - - private: - // We'll refer to the `LinkedGroup` type quite a bit below, so we shadow the class template - // with our own convenience typedef. But, we need to give LinkedGroup friend access to us first. - friend class LinkedGroup; - using LinkedGroup = LinkedGroup; - - protected: - // A proxy class representing this LinkedObject's LinkedGroup, if any. Instead of exposing - // multiple functions on LinkedObject, we expose one: `linkedGroup()`, and that function returns - // an object of this proxy class (or the similar ConstLinkedGroupProxy class below). - class LinkedGroupProxy { - public: - LinkedGroupProxy(LinkedObject& self): self(self) {} - void set(LinkedGroup& newGroup) { - self.setGroup(newGroup); - } - void set(kj::None) { - self.invalidateGroup(); - } - kj::Maybe tryGet() { - return self.tryGetGroup(); - } - - private: - LinkedObject& self; - }; - - // Const version of LinkedGroupProxy, exposing only `tryGet()`. - class ConstLinkedGroupProxy { - public: - ConstLinkedGroupProxy(const LinkedObject& self): self(self) {} - kj::Maybe tryGet() const { - return self.tryGetGroup(); - } - - private: - const LinkedObject& self; - }; - - // Provide access to this Object's LinkedGroup, if any. - LinkedGroupProxy linkedGroup() { - return *this; - } - ConstLinkedGroupProxy linkedGroup() const { - return *this; - } - - private: - void setGroup(LinkedGroup& newGroup) { - // Invalidate our current group membership, if any. - KJ_IF_SOME(oldGroup, maybeGroup) { - // If we're already a member of `newGroup`, we're done. Otherwise, we must remove ourselves - // from the old group. - if (&newGroup == &oldGroup) { - return; - } else { - removeFromGroup(oldGroup); - } - } else { - KJ_IREQUIRE(!link.isLinked()); - } - - // Add ourselves to the new group. - newGroup.list.add(*this); - maybeGroup = newGroup; - } - - kj::Maybe tryGetGroup() { - KJ_IF_SOME(group, maybeGroup) { - KJ_IREQUIRE(link.isLinked()); - return static_cast(group); - } else { - KJ_IREQUIRE(!link.isLinked()); - return kj::none; - } - } - - kj::Maybe tryGetGroup() const { - KJ_IF_SOME(group, maybeGroup) { - KJ_IREQUIRE(link.isLinked()); - return static_cast(group); - } else { - KJ_IREQUIRE(!link.isLinked()); - return kj::none; - } - } - - void invalidateGroup() { - KJ_IF_SOME(group, maybeGroup) { - removeFromGroup(group); - } else { - KJ_IREQUIRE(!link.isLinked()); - } - } - - // Helper for `setGroup()`, `invalidateGroup()`, and `~LinkedGroup()`. - void removeFromGroup(LinkedGroup& group) { - KJ_IREQUIRE(link.isLinked()); - group.list.remove(*this); - maybeGroup = kj::none; - } - - kj::ListLink link; - kj::Maybe maybeGroup; -}; - -// An iterator which wraps `InnerIterator` and `static_cast`s all mutable dereferences to -// `MaybeConstT&`, and all const dereferences to `const T&`. -// -// With the Ranges TS, all of this nonsense could be boiled down to a one-liner based on -// `std::views::transform()`. I encountered too many puzzles to solve while trying to get that -// working, so here we are. -template -class StaticCastIterator { - public: - // Construct an iterator using a default-constructed InnerIterator. In practice, this constructs - // an end iterator. - StaticCastIterator() = default; - - // Construct an iterator wrapping `inner`. - StaticCastIterator(InnerIterator inner): inner(inner) {} - - MaybeConstT& operator*() { - return static_cast(*inner); - } - const T& operator*() const { - return static_cast(*inner); - } - MaybeConstT* operator->() { - return static_cast(inner.operator->()); - } - const T* operator->() const { - return static_cast(inner.operator->()); - } - - inline StaticCastIterator& operator++() { - ++inner; - return *this; - } - inline StaticCastIterator operator++(int) { - StaticCastIterator result = *this; - ++inner; - return result; - } - - inline bool operator==(const StaticCastIterator& other) const { - return inner == other.inner; - } - - private: - InnerIterator inner; -}; - -} // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/maybe.rs b/src/rust/cxx/kj-rs/maybe.rs index 2b5b015c5ee..2a528d39aff 100644 --- a/src/rust/cxx/kj-rs/maybe.rs +++ b/src/rust/cxx/kj-rs/maybe.rs @@ -1,3 +1,9 @@ +//! FFI island: the `KjMaybe` representation of `kj::Maybe`. +//! +//! (See crate-root `#![deny(unsafe_code)]`.) Carries the `unsafe trait` niche contracts +//! (`HasNiche`/`MaybeItem`) and `assume_init` on the discriminated union. A genuine unsafe seam. +#![allow(unsafe_code)] + use std::mem::MaybeUninit; use std::pin::Pin; @@ -31,11 +37,12 @@ unsafe trait HasNiche: Sized { fn is_niche(value: *const Self) -> bool; } -// In Rust, references are not allowed to be null, so a null `MaybeUninit<&T>` is a niche -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// SAFETY: in Rust, references are not allowed to be null, so a null `MaybeUninit<&T>` is a +// niche (see the `HasNiche` trait contract above). unsafe impl HasNiche for &T { fn is_niche(value: *const &T) -> bool { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `value` points to a valid `&T`; we read it as a `*const *const T` (never as a + // reference, which the compiler assumes non-null) to test the pointer for null. unsafe { // We must cast it as pointing to a pointer, as opposed to a reference, // because the rust compiler assumes a reference is never null, and @@ -45,10 +52,10 @@ unsafe impl HasNiche for &T { } } -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// SAFETY: as for `&T` — a null `&mut T` is the niche (see the `HasNiche` trait contract). unsafe impl HasNiche for &mut T { fn is_niche(value: *const &mut T) -> bool { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `value` points to a valid `&mut T`; read as `*const *mut T` to null-check. unsafe { // We must cast it as pointing to a pointer, as opposed to a reference, // because the rust compiler assumes a reference is never null, and @@ -58,10 +65,11 @@ unsafe impl HasNiche for &mut T { } } -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// SAFETY: as for `&mut T` — a null pointee is the niche (see the `HasNiche` trait contract). unsafe impl HasNiche for Pin<&mut T> { fn is_niche(value: *const Pin<&mut T>) -> bool { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `value` points to a valid `Pin<&mut T>` (layout-identical to `&mut T`); read + // as `*const *mut T` to null-check. unsafe { // We must cast it as pointing to a pointer, as opposed to a reference, // because the rust compiler assumes a reference is never null, and @@ -72,10 +80,10 @@ unsafe impl HasNiche for Pin<&mut T> { } // In `kj`, `kj::Own` are considered `none` in a `Maybe` if the data pointer is null -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// SAFETY: a `KjOwn` with a null data pointer is `kj::none` (see the `HasNiche` trait contract). unsafe impl HasNiche for crate::repr::KjOwn { fn is_niche(value: *const Self) -> bool { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `value` points to a valid `KjOwn`; querying its data pointer is sound. unsafe { (*value).as_ptr().is_null() } } } @@ -111,7 +119,8 @@ pub unsafe trait MaybeItem: Sized { } fn drop_in_place(value: &mut KjMaybe) { if ::is_some(value) { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `is_some` just confirmed the `some` union member is initialized, so + // dropping it in place is sound. `KjMaybe`'s `Drop` calls this exactly once. unsafe { value.some.assume_init_drop(); } @@ -123,7 +132,9 @@ pub unsafe trait MaybeItem: Sized { /// Avoids running into generic specialization problems. macro_rules! impl_maybe_item_for_has_niche { ($ty:ty) => { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `$ty` is only ever a `HasNiche` type (enforced at the macro's use sites), so + // it carries a `()` discriminant and detects `none` via its null niche — matching kj's + // niche-value-optimized `Maybe` layout, as the `MaybeItem` trait contract requires. unsafe impl MaybeItem for $ty { type Discriminant = (); @@ -160,7 +171,9 @@ macro_rules! impl_maybe_item_for_has_niche { /// Avoids running into generic specialization problems. macro_rules! impl_maybe_item_for_primitive { ($ty:ty) => { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: primitives have no niche, so this mirrors kj's non-niche + // `kj::_::NullableValue` layout with an explicit `bool` discriminant (`is_set`) + // followed by the value, exactly as the `MaybeItem` trait contract requires. unsafe impl MaybeItem for $ty { type Discriminant = bool; @@ -198,7 +211,8 @@ impl_maybe_item_for_primitive!( u8, u16, u32, u64, u128, usize, i8, i16, i32, i64, i128, isize, f32, f64, bool, &str, String ); -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// SAFETY: `&[T]` is a fat pointer with no usable niche here, so it uses the explicit +// `bool`-discriminant (non-niche) `MaybeItem` representation, matching kj's layout. unsafe impl MaybeItem for &[T] { type Discriminant = bool; @@ -233,7 +247,9 @@ unsafe impl MaybeItem for &[T] { // // We therefore mirror that layout with a `bool` discriminant here, exactly like // the primitive types above, rather than implementing [`HasNiche`]. -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// +// SAFETY: `kj::Rc` defines no `Maybe` niche members, so `kj::Maybe>` uses the +// non-niche `bool`-discriminant `NullableValue` layout mirrored here (see comment above). unsafe impl MaybeItem for crate::KjRc { type Discriminant = bool; @@ -260,7 +276,9 @@ unsafe impl MaybeItem for crate::KjRc { } } -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. +// SAFETY: like `kj::Rc`, `kj::Arc` defines no `Maybe` niche members, so +// `kj::Maybe>` uses the non-niche `bool`-discriminant `NullableValue` layout +// mirrored here. unsafe impl MaybeItem for crate::KjArc { type Discriminant = bool; @@ -387,7 +405,9 @@ pub(crate) mod repr { if value.is_some() { // We can't move out of value so we copy it and forget it in // order to perform a "manual" move out of value - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `is_some` confirmed `some` is initialized; `assume_init_read` copies + // it out, and the immediately following `mem::forget(value)` prevents the + // source from being dropped, so ownership moves out exactly once. let ret = unsafe { Some(value.some.assume_init_read()) }; std::mem::forget(value); ret @@ -408,10 +428,10 @@ pub(crate) mod repr { if self.is_none() { write!(f, "Maybe::None") } else { - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - write!(f, "Maybe::Some({:?})", unsafe { - self.some.assume_init_ref() - }) + // SAFETY: the `is_none()` branch above is false here, so `some` is + // initialized and may be borrowed for formatting. + let value = unsafe { self.some.assume_init_ref() }; + write!(f, "Maybe::Some({value:?})") } } } diff --git a/src/rust/cxx/kj-rs/own.rs b/src/rust/cxx/kj-rs/own.rs index 2061fb60168..a89fc5a1f09 100644 --- a/src/rust/cxx/kj-rs/own.rs +++ b/src/rust/cxx/kj-rs/own.rs @@ -1,4 +1,8 @@ //! The `workerd-cxx` module containing the [`Own`] type, which is bindings to the `kj::Own` C++ type +//! +//! FFI island (see crate-root `#![deny(unsafe_code)]`): `KjOwn` mirrors `kj::Own` — raw-pointer +//! deref and `extern "C"` disposer/refcount calls. A genuine unsafe seam. +#![allow(unsafe_code)] use std::fmt; use std::marker::PhantomData; @@ -21,18 +25,16 @@ impl NonNullExceptMaybe { } pub unsafe fn as_ref(&self) -> &T { - // Safety: - // This value will only be null when in a [`Maybe`], which does niche value optimization - // for a null pointer, so the inner [`Own`] can never be accessed if it is null - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self.0` is null only when this `NonNullExceptMaybe` lives inside a + // `Maybe` (which niche-optimizes the null pointer and never dereferences the inner + // `Own`), so here — reached only through the non-null `Own` API — it is a valid, + // live pointer. The caller's `unsafe` obligation is that `self` outlives the borrow. unsafe { &*self.0 } } pub unsafe fn as_mut(&mut self) -> &mut T { - // Safety: - // This value will only be null when in a [`Maybe`], which does niche value optimization - // for a null pointer, so the inner [`Own`] can never be accessed if it is null - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: as in `as_ref`, `self.0` is non-null and live when reached through the + // `Own` API; `&mut self` gives exclusive access, so the mutable reborrow is unique. unsafe { &mut *self.0 } } } @@ -119,11 +121,19 @@ pub mod repr { } } - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe impl Send for KjOwn where T: Send {} - - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe impl Sync for KjOwn where T: Sync {} + // NO `Send`/`Sync` impls, deliberately. + // + // A `KjOwn` carries a type-erased `kj::Disposer*` alongside the object pointer, and + // dropping the `KjOwn` runs that disposer on whichever thread the drop happens on. A + // bound on `T` alone (e.g. `T: Send`) says nothing about the disposer: `kj::Own`s minted + // from `kj::Rc::toOwn()`/`kj::refcounted` (non-atomic refcount decrement), arena-backed + // objects, or any other custom disposer are NOT safe to destroy from another thread, and + // nothing at the bridge boundary guarantees disposer thread-safety. + // + // All current consumers keep `KjOwn`s on the KJ event-loop thread that created them, so + // no impls are needed. If a genuine cross-thread use case appears, it must come with an + // explicit opt-in mechanism that asserts the *disposer* is thread-safe (not just `T`); + // do not re-add blanket impls here. impl Deref for KjOwn { type Target = T; @@ -207,7 +217,8 @@ pub mod repr { } let this = std::ptr::from_mut::(self).cast::(); - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `this` points to this live `KjOwn` being dropped exactly once; the C++ + // `own$drop` shim invokes the type-erased `kj::Disposer` stored alongside `ptr`. unsafe { __drop(this); } diff --git a/src/rust/cxx/kj-rs/promise.c++ b/src/rust/cxx/kj-rs/promise.c++ index 4a03da29f1b..8382c8175a5 100644 --- a/src/rust/cxx/kj-rs/promise.c++ +++ b/src/rust/cxx/kj-rs/promise.c++ @@ -11,8 +11,8 @@ namespace kj_rs { static_assert(sizeof(OwnPromiseNode) == sizeof(uint64_t) * 1, "OwnPromiseNode size changed"); static_assert(alignof(OwnPromiseNode) == alignof(uint64_t) * 1, "OwnPromiseNode alignment changed"); -void own_promise_node_drop_in_place(OwnPromiseNode* node) { - kj::dtor(*node); +void own_promise_node_drop_in_place(OwnPromiseNode& node) { + kj::dtor(node); } } // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/promise.h b/src/rust/cxx/kj-rs/promise.h index b4d79d07ae9..fc5a858df4a 100644 --- a/src/rust/cxx/kj-rs/promise.h +++ b/src/rust/cxx/kj-rs/promise.h @@ -10,7 +10,7 @@ namespace kj_rs { using OwnPromiseNode = kj::_::OwnPromiseNode; -void own_promise_node_drop_in_place(OwnPromiseNode*); +void own_promise_node_drop_in_place(OwnPromiseNode&); namespace repr { diff --git a/src/rust/cxx/kj-rs/promise.rs b/src/rust/cxx/kj-rs/promise.rs index 1f221b8af3b..b7d6abda451 100644 --- a/src/rust/cxx/kj-rs/promise.rs +++ b/src/rust/cxx/kj-rs/promise.rs @@ -1,3 +1,8 @@ +//! FFI island (see crate-root `#![deny(unsafe_code)]`): `OwnPromiseNode`/`PromiseFuture` bridge — +//! `unsafe impl ExternType`, `unsafe extern "C"` unwrap callbacks, and `Pin` projection. A genuine +//! unsafe seam. +#![allow(unsafe_code)] + use std::ffi::c_void; use std::future::Future; use std::marker::PhantomData; @@ -20,16 +25,11 @@ pub struct OwnPromiseNode(*mut c_void /* kj::_::PromiseNode* */); // It is forgotten using `MaybeUninit` and its ownership passed over to c++ in `unwrap`. impl Drop for OwnPromiseNode { fn drop(&mut self) { - // Safety: - // 1. Pointer to self is non-null, and obviously points to valid memory. - // 2. We do not read or write to the OwnPromiseNode's memory, so there are no atomicity nor - // interleaved pointer/reference access concerns. - // - // https://doc.rust-lang.org/std/ptr/index.html#safety - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe { - crate::ffi::own_promise_node_drop_in_place(self); - } + // `own_promise_node_drop_in_place` placement-destructs the node behind `self`. The + // borrow is valid for the call; the value is only logically dead afterwards, inside + // this `drop`, and the inner `*mut c_void` has no drop glue, so there is no + // use-after-free or double-free. Expressed as a `&mut` binding, so no `unsafe` needed. + crate::ffi::own_promise_node_drop_in_place(self); } } @@ -141,12 +141,36 @@ impl KjPromise for CallbacksFuture { // unwrap will take over node ownership let node = ManuallyDrop::new(node); - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `node.0` is a live `OwnPromiseNode` whose ownership the callback takes over + // (wrapped in `ManuallyDrop` so we don't also drop it); `ret` is valid, suitably-aligned + // uninitialized storage for `Output`, which the callback initializes on the success path. unsafe { (callbacks.unwrap)(node.0, ret.as_mut_ptr().cast::()).into_result() }?; - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: the `?` above propagated any error, so on this path the callback reported + // success and therefore initialized `ret`. Ok(unsafe { ret.assume_init() }) } } -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. -unsafe impl Send for CallbacksFuture {} +// No `unsafe impl Send for CallbacksFuture`, deliberately. +// +// `CallbacksFuture` is only ever wrapped in `PromiseFuture`, whose `PromiseAwaiter` holds an +// `Option` (a raw pointer, hence `!Send`), so the composed future is `!Send` +// regardless. The bridged async machinery is confined to the KJ event-loop thread and `spawn` is +// `spawn_local`-backed (no `Send` requirement), so nothing needs a `Send` impl. Asserting the +// wrapper stays `!Send` locks that in. +#[cfg(test)] +mod send_guards { + use static_assertions::assert_not_impl_any; + + use super::CallbacksFuture; + use super::PromiseFuture; + + // The raw `*mut c_void` node makes this `!Send`/`!Sync` on its own; guard against a future + // hand-written impl silently introducing cross-thread transfer of a KJ promise node. + assert_not_impl_any!(CallbacksFuture: Send, Sync); + + // After its first poll, `PromiseFuture`'s embedded awaiter memory is self-referential and + // event-loop-linked (see `PromiseAwaiter::_pinned`); it must stay `!Unpin` so safe code + // cannot move it between polls (`&mut`-based awaits require `Unpin`). + assert_not_impl_any!(PromiseFuture>: Unpin); +} diff --git a/src/rust/cxx/kj-rs/refcount.rs b/src/rust/cxx/kj-rs/refcount.rs index 0489a75884d..ebe88bcb7e1 100644 --- a/src/rust/cxx/kj-rs/refcount.rs +++ b/src/rust/cxx/kj-rs/refcount.rs @@ -1,4 +1,8 @@ //! Module for both [`KjRc`] and [`KjArc`], since they're nearly identical types +//! +//! FFI island (see crate-root `#![deny(unsafe_code)]`): `KjRc`/`KjArc` mirror `kj::Rc`/`kj::Arc` — +//! `unsafe impl Send/Sync`, `extern "C"` refcount ops, and `Pin` projection. A genuine unsafe seam. +#![allow(unsafe_code)] use static_assertions::assert_eq_align; use static_assertions::assert_eq_size; @@ -30,10 +34,26 @@ pub mod repr { ptr: NonNull, } - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe impl Send for KjArc where T: Send {} - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe impl Sync for KjArc where T: Sync {} + // Safety: `KjArc` mirrors `std::sync::Arc`'s thread-safety contract, and therefore + // requires the same `T: Send + Sync` bound for both `Send` and `Sync`: + // + // - `T: Sync` is required because clones can be sent to other threads, giving multiple + // threads concurrent `&T` access to the same pointee. + // - `T: Send` is required because the last `KjArc` to drop destroys the pointee on + // whichever thread it happens to live on, effectively transferring ownership of `T` + // to that thread. (Likewise, `get_mut()` can hand out exclusive access on any thread.) + // + // The reference count itself is managed on the C++ side by `kj::AtomicRefcounted` + // (atomic increments/decrements; the bridge's clone/drop shims require the pointee to be + // atomic-refcounted), so concurrent clone/drop of separate handles is safe once `T` + // satisfies the bounds above. + // + // A weaker `Send where T: Send` bound would be unsound: with `T: Send + !Sync`, cloning and + // sending a clone yields concurrent `&T` on two threads, so both impls require `T: Send + Sync`. + unsafe impl Send for KjArc where T: Send + Sync {} + // SAFETY: see the `Send` impl above — `KjArc` mirrors `std::sync::Arc`'s `T: Send + Sync` + // contract for `Sync` for the same reasons. + unsafe impl Sync for KjArc where T: Send + Sync {} impl KjRc { #[must_use] @@ -43,7 +63,8 @@ pub mod repr { fn __is_shared(this: *const c_void) -> bool; } - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self` is a live `KjRc` (`&self`), so its `*const c_void` refcounted + // pointer is valid for the C++ `is_shared` query. unsafe { __is_shared(std::ptr::from_ref(self).cast::()) } } @@ -73,7 +94,8 @@ pub mod repr { fn __drop(this: *mut c_void); } - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self` is a live `KjRc` being dropped exactly once; the C++ `drop` shim + // releases its refcount handle. unsafe { __drop(std::ptr::from_mut(self).cast::()); } @@ -88,7 +110,8 @@ pub mod repr { fn __is_shared(this: *const c_void) -> bool; } - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self` is a live `KjArc` (`&self`), so its `*const c_void` refcounted + // pointer is valid for the C++ `is_shared` query. unsafe { __is_shared(std::ptr::from_ref(self).cast::()) } } @@ -140,7 +163,8 @@ pub mod repr { } let mut ret = std::mem::MaybeUninit::::uninit(); - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self` is a live `KjRc`; the C++ `clone` shim bumps the refcount and + // initializes `ret` with a valid `KjRc`, so `assume_init` is sound afterwards. unsafe { __clone( std::ptr::from_ref(self).cast::(), @@ -159,7 +183,8 @@ pub mod repr { } let mut ret = std::mem::MaybeUninit::::uninit(); - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self` is a live `KjArc`; the C++ `clone` shim bumps the atomic refcount + // and initializes `ret` with a valid `KjArc`, so `assume_init` is sound afterwards. unsafe { __clone( std::ptr::from_ref(self).cast::(), @@ -177,7 +202,8 @@ pub mod repr { fn __drop(this: *mut c_void); } - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. + // SAFETY: `self` is a live `KjArc` being dropped exactly once; the C++ `drop` shim + // releases its refcount handle. unsafe { __drop(std::ptr::from_mut(self).cast::()); } diff --git a/src/rust/cxx/kj-rs/tests/BUILD.bazel b/src/rust/cxx/kj-rs/tests/BUILD.bazel index 204d928f8f1..b513baa6bf7 100644 --- a/src/rust/cxx/kj-rs/tests/BUILD.bazel +++ b/src/rust/cxx/kj-rs/tests/BUILD.bazel @@ -31,6 +31,7 @@ rust_library( # TODO(cleanup): Why isn't :cxx transitive? "//src/rust/cxx", "//src/rust/cxx/kj-rs", + "@crates_vendor//:static_assertions", ], ) @@ -139,10 +140,31 @@ wd_cc_library( ) cc_test( - name = "linked-group-test", + name = "shared-event-test", size = "small", srcs = [ - "linked-group-test.c++", + "shared-event-test.c++", + ], + linkstatic = select({ + "@platforms//os:windows": True, + "//conditions:default": False, + }), + target_compatible_with = select({ + "@//build/config:no_build": ["@platforms//:incompatible"], + "//conditions:default": [], + }), + deps = [ + "//src/rust/cxx/kj-rs", + "//src/rust/cxx/third-party:runtime", + "@capnp-cpp//src/kj:kj-test", + ], +) + +cc_test( + name = "neutralize-waker-test", + size = "small", + srcs = [ + "neutralize-waker-test.c++", ], linkstatic = select({ "@platforms//os:windows": True, diff --git a/src/rust/cxx/kj-rs/tests/awaitables-cc-test.c++ b/src/rust/cxx/kj-rs/tests/awaitables-cc-test.c++ index 22e22776ad7..e54744e06de 100644 --- a/src/rust/cxx/kj-rs/tests/awaitables-cc-test.c++ +++ b/src/rust/cxx/kj-rs/tests/awaitables-cc-test.c++ @@ -8,6 +8,16 @@ #include +// Raw-RustFuture test helpers, defined in tests/lib.rs and tests/test_futures.rs. The +// bridge's generated `async fn` shims always apply RustFuture's eager-by-default +// kj::Promise conversion, so tests that need a *cold* promise receive the not-yet-converted +// RustFuture through these and call `.lazily()` (kj-rs/future.h) themselves. +extern "C" { +void kj_rs_demo_lazy_side_effect_future(::kj_rs::repr::RustFuture* out); +void kj_rs_demo_lazy_future_awaiting_cancellable_promise(::kj_rs::repr::RustFuture* out); +void kj_rs_demo_work_before_poll(uint64_t* target, ::kj_rs::repr::RustFuture* out); +} + namespace kj_rs_demo { namespace { @@ -47,15 +57,8 @@ KJ_TEST("c++ can receive synchronous wakes during poll()") { for (auto testCase: std::initializer_list{ {CloningAction::None, WakingAction::WakeByRefSameThread}, - {CloningAction::None, WakingAction::WakeByRefBackgroundThread}, {CloningAction::CloneSameThread, WakingAction::WakeByRefSameThread}, - {CloningAction::CloneSameThread, WakingAction::WakeByRefBackgroundThread}, - {CloningAction::CloneBackgroundThread, WakingAction::WakeByRefSameThread}, - {CloningAction::CloneBackgroundThread, WakingAction::WakeByRefBackgroundThread}, {CloningAction::CloneSameThread, WakingAction::WakeSameThread}, - {CloningAction::CloneSameThread, WakingAction::WakeBackgroundThread}, - {CloningAction::CloneBackgroundThread, WakingAction::WakeSameThread}, - {CloningAction::CloneBackgroundThread, WakingAction::WakeBackgroundThread}, {CloningAction::WakeByRefThenCloneSameThread, WakingAction::WakeSameThread}, }) { auto waking = new_waking_future_void(testCase.cloningAction, testCase.wakingAction); @@ -168,6 +171,57 @@ KJ_TEST(".awaiting a Promise from Rust can produce an Err Result") { waitScope); } +KJ_TEST("a panicking bridged async fn surfaces as a catchable kj::Exception, not an abort") { + // Unwind protection in the RustFuture vtable (kj-rs/future.rs): panics escaping poll() are + // converted into errored completions, mirroring the sync bridge's panic -> kj::Exception + // conversion. Before that protection, any of these would abort the process (unwinding out + // of an extern "C" fn). + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + { + // Fallible future, panic on first poll. + auto exception = KJ_ASSERT_NONNULL( + kj::runCatchingExceptions([&]() { new_panicking_future_void().wait(waitScope); })); + KJ_EXPECT(exception.getDescription().contains("bridged future panicked on purpose"), + exception.getDescription()); + } + + { + // Infallible future: the promise can still reject (kj::Promise always carries an + // exception channel even when the Rust signature is infallible). + auto exception = KJ_ASSERT_NONNULL(kj::runCatchingExceptions( + [&]() { new_panicking_infallible_future_void().wait(waitScope); })); + KJ_EXPECT(exception.getDescription().contains("bridged infallible future panicked on purpose"), + exception.getDescription()); + } + + { + // Panic after a suspension point: exercises the event-loop-driven poll path (not the + // eager creation-time poll). + auto exception = KJ_ASSERT_NONNULL(kj::runCatchingExceptions( + [&]() { new_panicking_after_await_future_void().wait(waitScope); })); + KJ_EXPECT(exception.getDescription().contains("panicked after a suspension point"), + exception.getDescription()); + } + + // A panicking bridged future can also be caught from a KJ coroutine. + []() -> kj::Promise { + kj::Maybe maybeException; + try { + co_await new_panicking_future_void(); + } catch (...) { + maybeException = kj::getCaughtExceptionAsKj(); + } + auto& exception = KJ_ASSERT_NONNULL(maybeException, "should have thrown"); + KJ_EXPECT(exception.getDescription().contains("bridged future panicked on purpose"), + exception.getDescription()); + }().wait(waitScope); + + // The loop is still healthy after the panics: run a normal future to completion. + []() -> kj::Promise { co_await new_ready_future_void(); }().wait(waitScope); +} + KJ_TEST("Rust can await Promise") { kj::EventLoop loop; kj::WaitScope waitScope(loop); @@ -188,9 +242,11 @@ KJ_TEST("C++ can receive asynchronous wakes after poll()") { kj::WaitScope waitScope(loop); auto promise = new_threaded_delay_future_void(); - // It's not ready yet. + // It's not ready yet: the future stashed a clone of its waker and returned Pending. KJ_EXPECT(!promise.poll(waitScope)); - // But later it is. + // Wake the stashed waker on the loop thread; this arms the FuturePollEvent so the next poll + // completes. Exercises a cloned-waker wake that arrives after poll() has already returned. + wake_delayed_future(); promise.wait(waitScope); } @@ -199,14 +255,71 @@ KJ_TEST("Work before poll") { kj::WaitScope waitScope(loop); uint64_t val = 0; - // It should be possible for rust function to do work before returning the future - // even if we don't poll or cancel it. - auto promise = work_before_poll(val); + // It should be possible for a Rust function to do work before returning the future + // even if we don't poll or cancel it. The future panics if polled, so it is converted + // with RustFuture::lazily() (the eager-by-default conversion polls at creation); this + // also proves cold promises really are never polled unawaited. + ::kj_rs::repr::RustFuture fut; + kj_rs_demo_work_before_poll(&val, &fut); + auto promise = fut.lazily(); KJ_EXPECT(val == 42); } +// ======================================================================================= +// Eager-by-default vs RustFuture::lazily(): bridged async fns surface as *eager* +// kj::Promises (polled to their first suspension at creation, like a KJ coroutine); +// `.lazily()` is the C++-side escape hatch that restores the cold future. + +KJ_TEST("bridged async fns are eager by default: the body runs at promise creation") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + reset_side_effect_counter(); + { + auto promise = new_side_effect_future_void(); + // No suspension points, so it ran to completion synchronously at creation, before any + // await or event-loop turn. + KJ_EXPECT(get_side_effect_counter() == 1); + } + KJ_EXPECT(get_side_effect_counter() == 1); +} + +KJ_TEST("RustFuture::lazily() opts out: nothing runs until the promise is first awaited") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + reset_side_effect_counter(); + ::kj_rs::repr::RustFuture fut; + kj_rs_demo_lazy_side_effect_future(&fut); + auto promise = fut.lazily(); + KJ_EXPECT(get_side_effect_counter() == 0); + + // Turning the event loop without awaiting the promise doesn't run it either. + kj::evalLater([]() {}).wait(waitScope); + KJ_EXPECT(get_side_effect_counter() == 0); + + promise.wait(waitScope); + KJ_EXPECT(get_side_effect_counter() == 1); +} + +KJ_TEST("eager promises still cancel on drop (never explicitly polled by the caller)") { + // Cancellation semantics are unchanged by eager evaluation: dropping the promise + // synchronously cancels the Rust future and the KJ promise it is awaiting. Here the + // caller never polls or awaits — creation alone started the future (suspending it at its + // .await), and destruction alone cancels it. + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + reset_cancellation_counter(); + { + auto promise = new_future_awaiting_cancellable_promise(); + KJ_EXPECT(get_cancellation_counter() == 0); + } + KJ_EXPECT(get_cancellation_counter() == 1); +} + // TODO(someday): More test cases. -// - Standalone ArcWaker tests. Ensure Rust calls ArcWaker destructor when we expect. +// - Standalone FutureWakerCell tests. Ensure Rust drops cloned waker cells when we expect. // - Throwing an exception from PromiseNode functions, including destructor. // ======================================================================================= @@ -219,11 +332,16 @@ KJ_TEST("Work before poll") { KJ_TEST("Cancellation: drop never-polled Rust future") { // Dropping a kj::Promise wrapping a Rust future that was never polled should not crash. Since the // future was never polled, the Rust async function body was never entered, so no sub-promises - // exist to cancel. + // exist to cancel. Uses a raw RustFuture converted with `.lazily()`: eager-by-default promises + // are always polled at least once (at creation), so only a cold promise can reach this path. kj::EventLoop loop; kj::WaitScope waitScope(loop); - { auto promise = new_future_awaiting_cancellable_promise(); } + { + ::kj_rs::repr::RustFuture fut; + kj_rs_demo_lazy_future_awaiting_cancellable_promise(&fut); + auto promise = fut.lazily(); + } } KJ_TEST("Cancellation: C++ dropping promise cancels Rust future's awaited KJ promise") { diff --git a/src/rust/cxx/kj-rs/tests/lib.rs b/src/rust/cxx/kj-rs/tests/lib.rs index 9a9aa7ae39d..31a253e2b12 100644 --- a/src/rust/cxx/kj-rs/tests/lib.rs +++ b/src/rust/cxx/kj-rs/tests/lib.rs @@ -13,6 +13,7 @@ mod test_own; mod test_refcount; use kj_rs::KjOwn; +use test_futures::get_side_effect_counter; use test_futures::new_drop_cancellable_promise_without_polling; use test_futures::new_error_handling_future_void_infallible; use test_futures::new_errored_future_void; @@ -20,17 +21,23 @@ use test_futures::new_future_awaiting_cancellable_promise; use test_futures::new_kj_errored_future_void; use test_futures::new_layered_ready_future_void; use test_futures::new_naive_select_future_void; +use test_futures::new_panicking_after_await_future_void; +use test_futures::new_panicking_future_void; +use test_futures::new_panicking_infallible_future_void; use test_futures::new_pending_future_void; use test_futures::new_promise_i32_awaiting_future_void; use test_futures::new_ready_future_i32; use test_futures::new_ready_future_void; use test_futures::new_select_with_cancellation; +use test_futures::new_side_effect_future_void; use test_futures::new_threaded_delay_future_void; use test_futures::new_two_step_cancellable_future; use test_futures::new_waking_future_void; use test_futures::new_wrapped_waker_future_void; use test_futures::poll_and_stash_promise_future; +use test_futures::reset_side_effect_counter; use test_futures::unstash_and_await_promise_future; +use test_futures::wake_delayed_future; use test_maybe::take_maybe_own; use test_maybe::take_maybe_own_ret; use test_maybe::take_maybe_ref; @@ -241,16 +248,13 @@ pub mod ffi { enum CloningAction { None, CloneSameThread, - CloneBackgroundThread, WakeByRefThenCloneSameThread, } enum WakingAction { None, WakeByRefSameThread, - WakeByRefBackgroundThread, WakeSameThread, - WakeBackgroundThread, } // Helper functions to create BoxFutureVoids for testing purposes. @@ -260,6 +264,9 @@ pub mod ffi { async fn new_ready_future_shared_type() -> Shared; async fn new_waking_future_void(cloning_action: CloningAction, waking_action: WakingAction); async fn new_threaded_delay_future_void(); + // Wakes the waker stashed by `new_threaded_delay_future_void`'s future, on the loop thread, + // to drive an asynchronous same-thread wake after poll() has returned. + fn wake_delayed_future(); async fn new_layered_ready_future_void() -> Result<()>; async fn new_naive_select_future_void() -> Result<()>; @@ -267,6 +274,12 @@ pub mod ffi { async fn new_errored_future_void() -> Result<()>; + // Unwind protection (kj-rs/future.rs): panics escaping a bridged future's poll() + // must become rejected promises (kj::Exception), not process aborts. + async fn new_panicking_future_void() -> Result<()>; + async fn new_panicking_infallible_future_void(); + async fn new_panicking_after_await_future_void() -> Result<()>; + async fn new_kj_errored_future_void() -> Result<()>; async fn new_error_handling_future_void_infallible(); @@ -275,7 +288,17 @@ pub mod ffi { async fn new_ready_future_i32(value: i32) -> Result; async fn new_pass_through_feature_shared() -> Shared; - async unsafe fn work_before_poll<'a>(target: &'a mut u64) -> Result<()>; + // Eager-by-default test helpers. The bridge's conversion polls the future + // synchronously to its first suspension at promise creation. The cold-promise + // (`RustFuture::lazily()`) counterparts bypass the bridge: they hand C++ the raw + // `RustFuture` through plain `extern "C"` helpers (see `test_futures.rs`). + #[expect(clippy::allow_attributes)] // Only called from C++ tests; #[expect(dead_code)] fails in builds where the lint does not fire + #[allow(dead_code)] + fn reset_side_effect_counter(); + #[expect(clippy::allow_attributes)] // Only called from C++ tests; #[expect(dead_code)] fails in builds where the lint does not fire + #[allow(dead_code)] + fn get_side_effect_counter() -> u64; + async fn new_side_effect_future_void(); // Cancellation test helpers. async fn new_future_awaiting_cancellable_promise() -> Result<()>; @@ -314,6 +337,20 @@ unsafe impl Send for ffi::OpaqueAtomicRefcountedClass {} // Safety: the test type follows the thread-safety contract of its C++ implementation. unsafe impl Sync for ffi::OpaqueAtomicRefcountedClass {} +// Compile-time thread-safety contracts (kj-rs/own.rs, kj-rs/refcount.rs): +// +// KjOwn is never Send/Sync: its type-erased kj disposer (kj::Rc, arena, ...) may not be +// thread-safe, regardless of T. +static_assertions::assert_not_impl_any!(KjOwn: Send, Sync); +static_assertions::assert_not_impl_any!(KjOwn: Send, Sync); +// KjArc matches std::sync::Arc: Send/Sync require T: Send + Sync... +static_assertions::assert_impl_all!(kj_rs::KjArc: Send, Sync); +// ...so a Send + !Sync payload (Cell) must make KjArc neither Send nor Sync (clones would +// otherwise hand concurrent &T to multiple threads). +static_assertions::assert_not_impl_any!(kj_rs::KjArc>: Send, Sync); +// KjRc (non-atomic refcount) must never be Send or Sync. +static_assertions::assert_not_impl_any!(kj_rs::KjRc: Send, Sync); + pub fn modify_own_return(mut own: KjOwn) -> KjOwn { own.pin_mut().set_data(72); own @@ -357,6 +394,33 @@ fn work_before_poll(target: &mut u64) -> impl Future> { } } +/// Hands C++ the raw, not-yet-converted future from [`work_before_poll`]. +/// +/// The returned future must never be polled (its body panics), so it cannot go through a +/// bridged `async fn` shim: those always apply `RustFuture`'s eager-by-default +/// `kj::Promise` conversion, which polls at creation. The C++ test converts it with +/// `RustFuture::lazily()` instead (see `awaitables-cc-test.c++`). +/// +/// # Safety +/// +/// `target` must be a valid, exclusive `u64` pointer that outlives the future; `out` must +/// point to uninitialized storage for one `::kj_rs::repr::RustFuture` (future.h), which the +/// caller takes ownership of. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn kj_rs_demo_work_before_poll<'a>( + target: &'a mut u64, + out: *mut kj_rs::repr::RustFuture<'a, ()>, +) { + let fut = kj_rs::repr::future(Box::pin(kj_rs::map_err( + work_before_poll(target), + file!(), + line!(), + ))); + // SAFETY: `out` points to uninitialized storage for one RustFuture, per this fn's + // `# Safety` contract. + unsafe { out.write(fut) }; +} + #[cfg(test)] mod tests { use crate::ffi; diff --git a/src/rust/cxx/kj-rs/tests/linked-group-test.c++ b/src/rust/cxx/kj-rs/tests/linked-group-test.c++ deleted file mode 100644 index 77a86ba2182..00000000000 --- a/src/rust/cxx/kj-rs/tests/linked-group-test.c++ +++ /dev/null @@ -1,364 +0,0 @@ -#include "kj-rs/linked-group.h" - -#include - -namespace kj_rs { -namespace { - -// Minimal concrete types for testing. -class TestGroup; -class TestObject; - -class TestGroup: public LinkedGroup { - public: - explicit TestGroup(int id): id(id) {} - - // Expose the protected member for testing. - using LinkedGroup::linkedObjects; - - int id; -}; - -class TestObject: public LinkedObject { - public: - explicit TestObject(int id): id(id) {} - - // Expose the protected member for testing. - using LinkedObject::linkedGroup; - - int id; -}; - -// --------------------------------------------------------------------------- -// Basic membership -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: object can join a group") { - TestGroup group(1); - TestObject object(10); - - KJ_EXPECT(group.linkedObjects().empty()); - KJ_EXPECT(object.linkedGroup().tryGet() == kj::none); - - object.linkedGroup().set(group); - - KJ_EXPECT(!group.linkedObjects().empty()); - KJ_EXPECT(group.linkedObjects().front().id == 10); - KJ_IF_SOME(g, object.linkedGroup().tryGet()) { - KJ_EXPECT(g.id == 1); - } else { - KJ_FAIL_EXPECT("expected group to be set"); - } -} - -KJ_TEST("LinkedGroup: object can leave a group") { - TestGroup group(1); - TestObject object(10); - - object.linkedGroup().set(group); - KJ_EXPECT(!group.linkedObjects().empty()); - - object.linkedGroup().set(kj::none); - KJ_EXPECT(group.linkedObjects().empty()); - KJ_EXPECT(object.linkedGroup().tryGet() == kj::none); -} - -KJ_TEST("LinkedGroup: object can switch groups") { - TestGroup group1(1); - TestGroup group2(2); - TestObject object(10); - - object.linkedGroup().set(group1); - KJ_EXPECT(!group1.linkedObjects().empty()); - KJ_EXPECT(group2.linkedObjects().empty()); - - object.linkedGroup().set(group2); - KJ_EXPECT(group1.linkedObjects().empty()); - KJ_EXPECT(!group2.linkedObjects().empty()); - KJ_IF_SOME(g, object.linkedGroup().tryGet()) { - KJ_EXPECT(g.id == 2); - } else { - KJ_FAIL_EXPECT("expected group to be set"); - } -} - -// --------------------------------------------------------------------------- -// Insertion order -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: objects are iterable in insertion order") { - TestGroup group(1); - TestObject a(1), b(2), c(3); - - a.linkedGroup().set(group); - b.linkedGroup().set(group); - c.linkedGroup().set(group); - - // Verify iteration order matches insertion order. - auto it = group.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it->id == 2); - ++it; - KJ_EXPECT(it->id == 3); - ++it; - KJ_EXPECT(it == group.linkedObjects().end()); -} - -// --------------------------------------------------------------------------- -// Redundant set() is a no-op -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: redundant set() does not change position") { - TestGroup group(1); - TestObject a(1), b(2), c(3); - - a.linkedGroup().set(group); - b.linkedGroup().set(group); - c.linkedGroup().set(group); - - // Re-set b to the same group — order should be unchanged. - b.linkedGroup().set(group); - - auto it = group.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it->id == 2); - ++it; - KJ_EXPECT(it->id == 3); - ++it; - KJ_EXPECT(it == group.linkedObjects().end()); -} - -// --------------------------------------------------------------------------- -// Lifetimes: object destroyed before group -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: destroying an object removes it from the group") { - TestGroup group(1); - TestObject a(1); - { - TestObject b(2); - b.linkedGroup().set(group); - a.linkedGroup().set(group); - - // Both present. - auto it = group.linkedObjects().begin(); - KJ_EXPECT(it->id == 2); - ++it; - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it == group.linkedObjects().end()); - } - // b is destroyed; only a remains. - auto it = group.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it == group.linkedObjects().end()); -} - -// --------------------------------------------------------------------------- -// Lifetimes: group destroyed before objects -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: destroying a group unlinks all objects") { - TestObject a(1), b(2); - { - TestGroup group(1); - a.linkedGroup().set(group); - b.linkedGroup().set(group); - KJ_EXPECT(a.linkedGroup().tryGet() != kj::none); - KJ_EXPECT(b.linkedGroup().tryGet() != kj::none); - } - // Group destroyed — objects should no longer reference it. - KJ_EXPECT(a.linkedGroup().tryGet() == kj::none); - KJ_EXPECT(b.linkedGroup().tryGet() == kj::none); -} - -// --------------------------------------------------------------------------- -// Iteration and removal of the front element -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: removing the front element during iteration is safe") { - TestGroup group(1); - TestObject a(1), b(2), c(3); - - a.linkedGroup().set(group); - b.linkedGroup().set(group); - c.linkedGroup().set(group); - - // The header documents that removing the *front* element during iteration is valid. - kj::Vector collected; - for (auto it = group.linkedObjects().begin(); it != group.linkedObjects().end();) { - auto& obj = *it; - ++it; // advance before removing - collected.add(obj.id); - obj.linkedGroup().set(kj::none); - } - - KJ_EXPECT(collected.size() == 3); - KJ_EXPECT(collected[0] == 1); - KJ_EXPECT(collected[1] == 2); - KJ_EXPECT(collected[2] == 3); - KJ_EXPECT(group.linkedObjects().empty()); -} - -// --------------------------------------------------------------------------- -// Const access -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: const access to group's objects") { - TestGroup group(1); - TestObject a(1), b(2); - - a.linkedGroup().set(group); - b.linkedGroup().set(group); - - const TestGroup& cgroup = group; - KJ_EXPECT(!cgroup.linkedObjects().empty()); - KJ_EXPECT(cgroup.linkedObjects().front().id == 1); - - auto it = cgroup.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it->id == 2); - ++it; - KJ_EXPECT(it == cgroup.linkedObjects().end()); -} - -KJ_TEST("LinkedGroup: const access to object's group") { - TestGroup group(1); - TestObject object(10); - - object.linkedGroup().set(group); - - const TestObject& cobject = object; - KJ_IF_SOME(g, cobject.linkedGroup().tryGet()) { - KJ_EXPECT(g.id == 1); - } else { - KJ_FAIL_EXPECT("expected group to be set"); - } -} - -// --------------------------------------------------------------------------- -// Empty state -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: default-constructed objects have no group") { - TestObject object(1); - KJ_EXPECT(object.linkedGroup().tryGet() == kj::none); -} - -KJ_TEST("LinkedGroup: default-constructed groups have no objects") { - TestGroup group(1); - KJ_EXPECT(group.linkedObjects().empty()); - KJ_EXPECT(group.linkedObjects().begin() == group.linkedObjects().end()); -} - -// --------------------------------------------------------------------------- -// Multiple objects, various removal patterns -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: removing a middle object leaves others intact") { - TestGroup group(1); - TestObject a(1), b(2), c(3); - - a.linkedGroup().set(group); - b.linkedGroup().set(group); - c.linkedGroup().set(group); - - b.linkedGroup().set(kj::none); - - auto it = group.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it->id == 3); - ++it; - KJ_EXPECT(it == group.linkedObjects().end()); -} - -KJ_TEST("LinkedGroup: removing all objects one by one") { - TestGroup group(1); - TestObject a(1), b(2), c(3); - - a.linkedGroup().set(group); - b.linkedGroup().set(group); - c.linkedGroup().set(group); - - a.linkedGroup().set(kj::none); - KJ_EXPECT(!group.linkedObjects().empty()); - - b.linkedGroup().set(kj::none); - KJ_EXPECT(!group.linkedObjects().empty()); - - c.linkedGroup().set(kj::none); - KJ_EXPECT(group.linkedObjects().empty()); -} - -// --------------------------------------------------------------------------- -// set(kj::none) on an unlinked object is a no-op -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: set(none) on an unlinked object is safe") { - TestObject object(1); - // Should not crash or assert. - object.linkedGroup().set(kj::none); - KJ_EXPECT(object.linkedGroup().tryGet() == kj::none); -} - -// --------------------------------------------------------------------------- -// Multiple groups are independent -// --------------------------------------------------------------------------- - -KJ_TEST("LinkedGroup: multiple groups are independent") { - TestGroup g1(1), g2(2); - TestObject a(1), b(2), c(3), d(4); - - a.linkedGroup().set(g1); - b.linkedGroup().set(g1); - c.linkedGroup().set(g2); - d.linkedGroup().set(g2); - - // Verify g1 has {a, b}. - { - auto it = g1.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it->id == 2); - ++it; - KJ_EXPECT(it == g1.linkedObjects().end()); - } - - // Verify g2 has {c, d}. - { - auto it = g2.linkedObjects().begin(); - KJ_EXPECT(it->id == 3); - ++it; - KJ_EXPECT(it->id == 4); - ++it; - KJ_EXPECT(it == g2.linkedObjects().end()); - } - - // Move b from g1 to g2. - b.linkedGroup().set(g2); - - { - auto it = g1.linkedObjects().begin(); - KJ_EXPECT(it->id == 1); - ++it; - KJ_EXPECT(it == g1.linkedObjects().end()); - } - { - auto it = g2.linkedObjects().begin(); - KJ_EXPECT(it->id == 3); - ++it; - KJ_EXPECT(it->id == 4); - ++it; - KJ_EXPECT(it->id == 2); - ++it; - KJ_EXPECT(it == g2.linkedObjects().end()); - } -} - -} // namespace -} // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/tests/neutralize-waker-test.c++ b/src/rust/cxx/kj-rs/tests/neutralize-waker-test.c++ new file mode 100644 index 00000000000..210a3e09eeb --- /dev/null +++ b/src/rust/cxx/kj-rs/tests/neutralize-waker-test.c++ @@ -0,0 +1,144 @@ +// Regression test: FutureWakerCell "neutralize-on-drop". +// +// The bridge's same-thread waker is a FutureWakerCell whose wake() arms the owning FuturePollEvent. +// The hazard: a waker clone is retained (e.g. handed to some sub-future) and its wake() fires AFTER +// the FuturePollEvent (and the boxed future it owns) has been torn down -- arming a freed +// kj::_::Event would be a use-after-free. +// +// This guards the refcounted-cell handling: the cell holds an Event*; the FuturePollEvent holds +// the strong ref and NULLS the cell on destruction (BEFORE the boxed future / sub-wakers drop). A +// retained clone that calls wake() after teardown observes null and is a safe no-op. +// +// Run under ASAN to confirm no UAF: +// bazel test //kj-rs/tests:neutralize-waker-test --config=asan + +#include +#include +#include +#include + +namespace kj_rs { +namespace { + +using kj::uint; + +// The refcounted cell shared between the FutureEvent and any retained Waker clones. The bridge is +// single-threaded (no cross-thread wakes), so a plain (non-atomic) kj::Refcounted with a bare +// Event* is sufficient -- no mutex/atomic needed. +class WakerCell: public kj::Refcounted { + public: + kj::_::Event* event = nullptr; + + bool observedNullOnWake = false; + uint wakeArmCount = 0; + + // Called by the FutureEvent's destructor to neutralize all outstanding waker clones. + void neutralize() { + event = nullptr; + } + + // The Waker's wake(): arm the FutureEvent, or no-op if it's been neutralized. + void wake() { + if (event != nullptr) { + event->armDepthFirst(); + ++wakeArmCount; + } else { + observedNullOnWake = true; // SAFE no-op: no arm of a freed Event, no UAF. + } + } +}; + +// Models the boxed Rust future (and its sub-wakers) owned inline by the FutureEvent. Its whole job +// here is to ASSERT the ordering requirement: by the time it is destroyed, the cell must already +// have been neutralized -- i.e. nulled BEFORE the boxed future / sub-wakers drop. +class BoxedFutureStandin { + public: + explicit BoxedFutureStandin(WakerCell& cell): cell(cell) {} + ~BoxedFutureStandin() noexcept(false) { + KJ_ASSERT(cell.event == nullptr, + "ordering violation: cell must be neutralized BEFORE the boxed future drops"); + } + + private: + WakerCell& cell; +}; + +// Stand-in for the FutureEvent: a kj Event that owns the WakerCell strong ref and the boxed future. +class FutureEventStandin final: public kj::_::Event { + public: + explicit FutureEventStandin(kj::Rc cellParam, kj::SourceLocation location = {}) + : Event(location), + cell(kj::mv(cellParam)), + boxedFuture(*cell) { + cell->event = this; + } + + ~FutureEventStandin() noexcept(false) { + // ORDERING: neutralize the cell FIRST (destructor body runs before member subobjects are + // destroyed). Member destruction order is reverse-declaration: `boxedFuture` then `cell`. + // So when boxedFuture's dtor asserts, the cell is already nulled. + cell->neutralize(); + } + + uint fireCount = 0; + void traceEvent(kj::_::TraceBuilder&) override {} + + private: + void fire() override { + ++fireCount; + } + + kj::Rc cell; // declared first -> destroyed LAST + BoxedFutureStandin boxedFuture; // declared second -> destroyed FIRST +}; + +KJ_TEST("FutureWaker neutralize-on-drop: retained clone wake() is a safe no-op after teardown") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + auto cell = kj::rc(); + auto retainedClone = cell.addRef(); // a second handle that will OUTLIVE the FutureEvent. + + auto event = kj::heap(kj::mv(cell)); + + // Sanity: while the FutureEvent is alive, a wake via the retained clone arms it. + retainedClone->wake(); + waitScope.poll(); + KJ_EXPECT(event->fireCount == 1); + KJ_EXPECT(retainedClone->wakeArmCount == 1); + KJ_EXPECT(!retainedClone->observedNullOnWake); + + // TEARDOWN: destroy the FutureEvent. Its dtor neutralizes the cell (asserted to happen before + // boxedFuture drops). The retained clone keeps the cell object itself alive. + event = nullptr; + + // The retained clone's wake() now observes a null Event* -> SAFE no-op. Without neutralize-on- + // drop this would arm a freed Event (UAF -- caught by ASAN under --config=asan). + retainedClone->wake(); + waitScope.poll(); // must not fire anything, must not crash + KJ_EXPECT(retainedClone->observedNullOnWake); + KJ_EXPECT(retainedClone->wakeArmCount == 1); // unchanged: the post-teardown wake armed nothing +} + +KJ_TEST("FutureWaker neutralize-on-drop: multiple retained clones all neutralized together") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + auto cell = kj::rc(); + auto cloneA = cell.addRef(); + auto cloneB = cell.addRef(); + + auto event = kj::heap(kj::mv(cell)); + event = nullptr; // teardown + + // Both retained clones observe null; neither arms a freed Event. + cloneA->wake(); + cloneB->wake(); + waitScope.poll(); + KJ_EXPECT(cloneA->observedNullOnWake); + KJ_EXPECT(cloneB->observedNullOnWake); + KJ_EXPECT(cloneA->wakeArmCount == 0); +} + +} // namespace +} // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/tests/shared-event-test.c++ b/src/rust/cxx/kj-rs/tests/shared-event-test.c++ new file mode 100644 index 00000000000..15da02bbfbf --- /dev/null +++ b/src/rust/cxx/kj-rs/tests/shared-event-test.c++ @@ -0,0 +1,321 @@ +// Regression test: one shared kj Event as the onReady target of MANY concurrent pending nodes. +// +// The bridge registers the SAME FuturePollEvent as the onReady target of every kj::Promise a Rust +// future is `.await`ing (many nodes -> one event, kj's native mechanism). This guards the +// kj Event/onReady properties that relies on: +// (a) idempotent arming: fulfilling several nodes in the SAME turn arms the one event once +// (Event::armDepthFirst's `if (prev == nullptr)` guard, async.c++:2201), +// (b) re-poll: fire() re-polls, ready nodes are consumed, not-yet-ready nodes stay registered +// and re-arm the event when THEY later resolve, +// (c) arm-while-firing: fulfilling a node DURING the shared event's own fire() safely re-arms +// it for the next turn (no "Promise callback destroyed itself" abort async.c++:2188, no +// lost wake) -- because turn() unlinks the event before fire(), so prev==nullptr during +// fire and armDepthFirst re-inserts it. +// +// Two test groups: +// GROUP A -- "pure kj primitive": N PromiseNodes each call node->onReady(&sharedEvent) DIRECTLY. +// Proves kj's OnReadyEvent::arm() -> Event::armDepthFirst() coalescing on one shared +// target. Readiness bookkeeping is test-driven (kj exposes no per-node readiness +// query): the point of Group A is the wake/arm coalescing path. +// GROUP B -- "faithful future model": each leaf is a trivial ~6-line per-leaf arm Event whose +// fire() sets ready=true and arms ONE shared re-poll event, and readiness here is +// genuinely kj-detected. + +#include +#include +#include +#include +#include + +namespace kj_rs { +namespace { + +using kj::uint; +using kj::_::Event; +using kj::_::ExceptionOr; +using kj::_::OwnPromiseNode; +using kj::_::PromiseNode; +using kj::_::Void; + +// A stable slot holding a fulfiller/node pair so we can call setSelfPointer() and later get(). +struct Slot { + kj::Own> fulfiller; + OwnPromiseNode node; + bool fulfilled = false; // test-side bookkeeping (Group A) + bool consumed = false; + + static kj::Own make() { + auto paf = kj::newPromiseAndFulfiller(); + auto self = kj::heap(); + self->fulfiller = kj::mv(paf.fulfiller); + self->node = PromiseNode::from(kj::mv(paf.promise)); + self->node->setSelfPointer(&self->node); + return self; + } + + void consume() { + ExceptionOr output; + node->get(output); + KJ_ASSERT(output.exception == kj::none); + consumed = true; + } +}; + +// ======================================================================================= +// GROUP A -- N nodes -> ONE shared event as direct onReady target. + +class SharedRepollEvent final: public Event { + public: + SharedRepollEvent(kj::ArrayPtr> slots, kj::SourceLocation location = {}) + : Event(location), + slots(slots) {} + + uint fireCount = 0; + uint consumedCount = 0; + + // If set, invoked once during the NEXT fire() -- models a sub-future resolving mid-poll. + kj::Function* armWhileFiringHook = nullptr; + + void traceEvent(kj::_::TraceBuilder&) override {} + + private: + kj::ArrayPtr> slots; + + void fire() override { + ++fireCount; + + // Model "the rust future re-polls all its sub-futures": consume every ready (fulfilled), + // not-yet-consumed node; leave the rest registered. + for (auto& slot: slots) { + if (slot->fulfilled && !slot->consumed) { + slot->consume(); + ++consumedCount; + } + } + + // ARM-WHILE-FIRING: fulfill another node during our own fire(). Its onReady points at us; + // arming us here must be safe (we were unlinked before fire, so prev==nullptr -> re-inserts). + if (armWhileFiringHook != nullptr) { + auto* hook = armWhileFiringHook; + armWhileFiringHook = nullptr; + (*hook)(); + } + // NOTE: we intentionally do NOT self-destruct; a FutureEvent lives until its future resolves. + } +}; + +KJ_TEST("SharedEvent(A): fulfilling several nodes in one turn arms the shared event idempotently") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + auto slots = kj::heapArray>(4); + for (auto& s: slots) s = Slot::make(); + + SharedRepollEvent shared(slots); + for (auto& s: slots) s->node->onReady(&shared); + + // Fulfill 3 of the 4 in the SAME turn (no loop run in between). + for (uint i: {0u, 1u, 2u}) { + slots[i]->fulfiller->fulfill(); + slots[i]->fulfilled = true; + } + + // Exactly one fire should happen: the 3 same-turn arms coalesced into a single armed event. + // (If they had NOT coalesced, fireCount would be 3.) + waitScope.poll(); + KJ_EXPECT(shared.fireCount == 1); + KJ_EXPECT(shared.consumedCount == 3); + + // (b) The 4th node was never fulfilled: it stays registered on the shared event. Fulfilling it + // now must re-arm the (idle) shared event and fire again. + slots[3]->fulfiller->fulfill(); + slots[3]->fulfilled = true; + waitScope.poll(); + KJ_EXPECT(shared.fireCount == 2); + KJ_EXPECT(shared.consumedCount == 4); +} + +KJ_TEST("SharedEvent(A): arm-while-firing re-arms safely for the next turn (no lost wake)") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + auto slots = kj::heapArray>(3); + for (auto& s: slots) s = Slot::make(); + + SharedRepollEvent shared(slots); + for (auto& s: slots) s->node->onReady(&shared); + + // During the shared event's FIRST fire(), fulfill slot[2]. Its onReady == &shared, so this arms + // `shared` while `shared` is mid-fire. Must not abort; must re-arm for the next turn. + kj::Function hook = [&]() { + slots[2]->fulfiller->fulfill(); + slots[2]->fulfilled = true; + }; + shared.armWhileFiringHook = &hook; + + // Kick off: fulfill slots 0 and 1 in this turn. + for (uint i: {0u, 1u}) { + slots[i]->fulfiller->fulfill(); + slots[i]->fulfilled = true; + } + + // A single poll() drains: fire #1 consumes 0,1 and (via the hook) fulfills slot[2], which arms + // `shared` mid-fire; fire #2 (the re-arm) consumes slot[2]. Proves no abort + no lost wake. + waitScope.poll(); + KJ_EXPECT(shared.fireCount == 2); + KJ_EXPECT(shared.consumedCount == 3); +} + +// ======================================================================================= +// GROUP B -- faithful future model: trivial per-leaf arm events + ONE shared re-poll event. + +class FutureEvent; + +// A leaf `.await` of a kj::Promise. ~6 lines of real logic: on its promise's readiness, record +// ready and arm the shared FutureEvent. +class LeafAwaiter final: public Event { + public: + LeafAwaiter(OwnPromiseNode nodeParam, FutureEvent& futureEvent, kj::SourceLocation location = {}); + ~LeafAwaiter() noexcept(false) { + node = nullptr; + } + + bool ready = false; + bool consumed = false; + + void consume() { + KJ_ASSERT(ready && !consumed); + ExceptionOr output; + node->get(output); + KJ_ASSERT(output.exception == kj::none); + consumed = true; + } + + void traceEvent(kj::_::TraceBuilder&) override {} + + private: + OwnPromiseNode node; + FutureEvent& futureEvent; + void fire() override; // defined after FutureEvent +}; + +// The ONE event that "is the future": its fire() re-polls the whole future (== consume any ready +// leaves, leave the rest). Every leaf arms THIS single event. +class FutureEvent final: public Event { + public: + FutureEvent(kj::SourceLocation location = {}): Event(location) {} + + uint fireCount = 0; + + // If set, invoked once during the NEXT fire() -- models a sub-future resolving mid-poll. + kj::Function* armWhileFiringHook = nullptr; + + void addLeaf(kj::Own leaf) { + leaves.add(kj::mv(leaf)); + } + + bool allConsumed() const { + for (auto& l: leaves) + if (!l->consumed) return false; + return true; + } + uint consumedCount() const { + uint n = 0; + for (auto& l: leaves) + if (l->consumed) ++n; + return n; + } + + void traceEvent(kj::_::TraceBuilder&) override {} + + private: + kj::Vector> leaves; + + void fire() override { + ++fireCount; + for (auto& l: leaves) { + if (l->ready && !l->consumed) l->consume(); + } + if (armWhileFiringHook != nullptr) { + auto* hook = armWhileFiringHook; + armWhileFiringHook = nullptr; + (*hook)(); // fulfills another leaf -> its LeafAwaiter arms `this` mid-fire + } + } +}; + +LeafAwaiter::LeafAwaiter(OwnPromiseNode nodeParam, FutureEvent& fe, kj::SourceLocation location) + : Event(location), + node(kj::mv(nodeParam)), + futureEvent(fe) { + node->setSelfPointer(&node); + node->onReady(this); +} + +void LeafAwaiter::fire() { + ready = true; + futureEvent.armDepthFirst(); // arm the ONE shared future event (idempotent across leaves) +} + +KJ_TEST( + "SharedEvent(B): many leaves arm one FutureEvent; ready consumed, pending stay registered") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + FutureEvent future; + kj::Vector>> fulfillers; + + for (uint i = 0; i < 4; ++i) { + auto paf = kj::newPromiseAndFulfiller(); + fulfillers.add(kj::mv(paf.fulfiller)); + future.addLeaf(kj::heap(PromiseNode::from(kj::mv(paf.promise)), future)); + } + + // Fulfill 3 leaves in one turn -> 3 leaf Events fire (each arms `future`), then `future` fires + // ONCE (coalesced). turn() runs the leaf events + the single future event. + for (uint i: {0u, 1u, 2u}) fulfillers[i]->fulfill(); + + loop.run(64); + KJ_EXPECT(future.fireCount >= 1); + KJ_EXPECT(future.consumedCount() == 3); + KJ_EXPECT(!future.allConsumed()); + + uint fireCountAfter3 = future.fireCount; + + // 4th leaf stays registered; fulfilling it re-arms the future event. + fulfillers[3]->fulfill(); + loop.run(64); + KJ_EXPECT(future.fireCount > fireCountAfter3); + KJ_EXPECT(future.allConsumed()); +} + +KJ_TEST( + "SharedEvent(B): leaf resolving during the future's fire re-arms the future (no lost wake)") { + kj::EventLoop loop; + kj::WaitScope waitScope(loop); + + FutureEvent future; + kj::Vector>> fulfillers; + + for (uint i = 0; i < 3; ++i) { + auto paf = kj::newPromiseAndFulfiller(); + fulfillers.add(kj::mv(paf.fulfiller)); + future.addLeaf(kj::heap(PromiseNode::from(kj::mv(paf.promise)), future)); + } + + // During the future's FIRST fire(), fulfill leaf 2. Its LeafAwaiter event will arm `future` + // while `future` is still mid-fire -> must safely re-arm for the next turn. + kj::Function hook = [&]() { fulfillers[2]->fulfill(); }; + future.armWhileFiringHook = &hook; + + // Kick off with leaves 0 and 1. + fulfillers[0]->fulfill(); + fulfillers[1]->fulfill(); + + loop.run(64); + KJ_EXPECT(future.fireCount >= 2); // at least: fire consuming 0/1, then the re-armed fire for 2 + KJ_EXPECT(future.allConsumed()); // leaf 2 (fulfilled during a fire) was not lost +} + +} // namespace +} // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/tests/test_futures.rs b/src/rust/cxx/kj-rs/tests/test_futures.rs index cb8213e6eba..9c0acb2f8c0 100644 --- a/src/rust/cxx/kj-rs/tests/test_futures.rs +++ b/src/rust/cxx/kj-rs/tests/test_futures.rs @@ -25,6 +25,54 @@ pub async fn new_ready_future_void() { std::future::ready(()).await } +// Eager-by-default vs `RustFuture::lazily()` test helpers: they record when their body ran +// so the C++ driver can observe eager promises running at creation and `.lazily()` ones only +// when awaited. The bridge's generated shims always apply the eager conversion, so the +// `.lazily()` tests receive the raw `RustFuture` through the plain `extern "C"` helpers +// below instead of bridged `async fn` declarations. + +static SIDE_EFFECT_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + +pub fn reset_side_effect_counter() { + SIDE_EFFECT_COUNTER.store(0, std::sync::atomic::Ordering::SeqCst); +} + +pub fn get_side_effect_counter() -> u64 { + SIDE_EFFECT_COUNTER.load(std::sync::atomic::Ordering::SeqCst) +} + +/// Increments the side-effect counter when its body runs (first poll), then completes. +pub async fn new_side_effect_future_void() { + SIDE_EFFECT_COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst); +} + +/// Same body as [`new_side_effect_future_void`], but handed to C++ as a raw `RustFuture` +/// (via [`kj_rs_demo_lazy_side_effect_future`]) and converted with `RustFuture::lazily()`: +/// the C++ promise is cold, so the counter only moves once the promise is first awaited. +pub async fn new_lazy_side_effect_future_void() { + SIDE_EFFECT_COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst); +} + +/// Hands C++ the raw, not-yet-converted [`new_lazy_side_effect_future_void`] future. +/// +/// `awaitables-cc-test.c++` converts it with `RustFuture::lazily()` (future.h). Plain +/// `extern "C"` because the bridge's generated `async fn` shims always apply the +/// eager-by-default `kj::Promise` conversion before C++ ever sees the future. +/// +/// # Safety +/// +/// `out` must point to uninitialized storage for one `::kj_rs::repr::RustFuture` +/// (future.h), which the caller takes ownership of. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn kj_rs_demo_lazy_side_effect_future( + out: *mut kj_rs::repr::RustInfallibleFuture<'static, ()>, +) { + let fut = kj_rs::repr::infallible_future(Box::pin(new_lazy_side_effect_future_void())); + // SAFETY: `out` points to uninitialized storage for one RustFuture, per this fn's + // `# Safety` contract. + unsafe { out.write(fut) }; +} + struct WakingFuture { done: bool, cloning_action: CloningAction, @@ -45,8 +93,7 @@ fn do_no_clone_wake(waker: &Waker, waking_action: WakingAction) { match waking_action { WakingAction::None => {} WakingAction::WakeByRefSameThread => waker.wake_by_ref(), - WakingAction::WakeByRefBackgroundThread => on_background_thread(|| waker.wake_by_ref()), - WakingAction::WakeSameThread | WakingAction::WakeBackgroundThread => { + WakingAction::WakeSameThread => { panic!("cannot wake() without cloning"); } _ => panic!("invalid WakingAction"), @@ -57,9 +104,7 @@ fn do_cloned_wake(waker: Waker, waking_action: WakingAction) { match waking_action { WakingAction::None => {} WakingAction::WakeByRefSameThread => waker.wake_by_ref(), - WakingAction::WakeByRefBackgroundThread => on_background_thread(|| waker.wake_by_ref()), WakingAction::WakeSameThread => waker.wake(), - WakingAction::WakeBackgroundThread => on_background_thread(move || waker.wake()), _ => panic!("invalid WakingAction"), } } @@ -81,10 +126,6 @@ impl Future for WakingFuture { let waker = waker.clone(); do_cloned_wake(waker, self.waking_action); } - CloningAction::CloneBackgroundThread => { - let waker = on_background_thread(|| waker.clone()); - do_cloned_wake(waker, self.waking_action); - } CloningAction::WakeByRefThenCloneSameThread => { waker.wake_by_ref(); let waker = waker.clone(); @@ -102,44 +143,52 @@ pub async fn new_waking_future_void(cloning_action: CloningAction, waking_action WakingFuture::new(cloning_action, waking_action).await } -struct ThreadedDelayFuture { - handle: Option>, +// A future that, on its first poll, stashes a clone of its waker and returns Pending WITHOUT +// waking; a later call to `wake_delayed_future()` — made by the C++ driver on the event loop's own +// thread, after poll() has already returned — wakes the stashed clone, arming the FuturePollEvent +// so the next poll completes. This exercises an *asynchronous* wake (one that arrives after poll() +// returned, via a cloned waker) entirely on the loop thread, which is all the single-thread bridge +// supports. + +thread_local! { + static DELAYED_WAKER: std::cell::RefCell> = const { std::cell::RefCell::new(None) }; } -impl ThreadedDelayFuture { - fn new() -> Self { - Self { handle: None } - } +struct DelayedWakeFuture { + done: bool, } -/// Run a function, `f`, on a thread in the background and return its result. -fn on_background_thread(f: impl FnOnce() -> T + Send) -> T { - std::thread::scope(|scope| match scope.spawn(f).join() { - Ok(value) => value, - Err(payload) => std::panic::resume_unwind(payload), - }) +impl DelayedWakeFuture { + fn new() -> Self { + Self { done: false } + } } -impl Future for ThreadedDelayFuture { +impl Future for DelayedWakeFuture { type Output = (); fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context) -> Poll<()> { - if let Some(handle) = self.handle.take() { - let _ = handle.join(); + if self.done { return Poll::Ready(()); } - - let waker = cx.waker(); - let waker = on_background_thread(|| waker.clone()); - self.handle = Some(std::thread::spawn(|| { - std::thread::sleep(std::time::Duration::from_millis(100)); - waker.wake(); - })); + // Stash a clone of the waker for the C++ driver to wake later, on this same thread. + DELAYED_WAKER.with(|w| *w.borrow_mut() = Some(cx.waker().clone())); + self.done = true; Poll::Pending } } pub async fn new_threaded_delay_future_void() { - ThreadedDelayFuture::new().await + DelayedWakeFuture::new().await +} + +/// Wake the waker stashed by [`DelayedWakeFuture`]. Called by the C++ test driver on the event +/// loop thread, after the future's first poll has returned Pending. +pub fn wake_delayed_future() { + DELAYED_WAKER.with(|w| { + if let Some(waker) = w.borrow_mut().take() { + waker.wake(); + } + }); } pub async fn new_layered_ready_future_void() -> Result<()> { @@ -214,6 +263,26 @@ pub async fn new_errored_future_void() -> Result<()> { Err(std::io::Error::other("test error")) } +// Unwind-protection helpers (kj-rs/future.rs vtable): a panic escaping a bridged future's +// poll() must surface to C++ as a rejected kj::Promise carrying a kj::Exception, not a +// process abort. These panic at different points to cover both the fallible and infallible +// vtables and both the first-poll and event-loop-driven poll paths. + +pub async fn new_panicking_future_void() -> Result<()> { + panic!("bridged future panicked on purpose"); +} + +pub async fn new_panicking_infallible_future_void() { + panic!("bridged infallible future panicked on purpose"); +} + +pub async fn new_panicking_after_await_future_void() -> Result<()> { + crate::ffi::new_ready_promise_void() + .await + .expect("should not throw"); + panic!("bridged future panicked after a suspension point"); +} + pub async fn new_kj_errored_future_void() -> std::result::Result<(), cxx::KjError> { Err(cxx::KjError::new( cxx::KjExceptionType::Overloaded, @@ -258,6 +327,43 @@ pub async fn new_future_awaiting_cancellable_promise() -> Result<()> { Ok(()) } +/// Like [`new_future_awaiting_cancellable_promise`], but handed to C++ as a raw +/// `RustFuture` (via [`kj_rs_demo_lazy_future_awaiting_cancellable_promise`]) and converted +/// with `RustFuture::lazily()`, so the C++ promise really is never polled unless awaited — +/// the only way to exercise the "drop a never-polled future" path now that bridged promises +/// are eager by default. +pub async fn new_lazy_future_awaiting_cancellable_promise() -> Result<()> { + crate::ffi::new_cancellation_detecting_promise_void() + .await + .map_err(Error::other)?; + Ok(()) +} + +/// Hands C++ the raw, not-yet-converted [`new_lazy_future_awaiting_cancellable_promise`] +/// future. +/// +/// `awaitables-cc-test.c++` converts it with `RustFuture::lazily()` (future.h). Plain +/// `extern "C"` because the bridge's generated `async fn` shims always apply the +/// eager-by-default `kj::Promise` conversion before C++ ever sees the future. +/// +/// # Safety +/// +/// `out` must point to uninitialized storage for one `::kj_rs::repr::RustFuture` +/// (future.h), which the caller takes ownership of. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn kj_rs_demo_lazy_future_awaiting_cancellable_promise( + out: *mut kj_rs::repr::RustFuture<'static, ()>, +) { + let fut = kj_rs::repr::future(Box::pin(kj_rs::map_err( + new_lazy_future_awaiting_cancellable_promise(), + file!(), + line!(), + ))); + // SAFETY: `out` points to uninitialized storage for one RustFuture, per this fn's + // `# Safety` contract. + unsafe { out.write(fut) }; +} + /// Two-step future: the first step completes normally, and the second step awaits a /// cancellation-detecting promise that never resolves. After one poll, the future will have /// advanced past step 1 and be suspended at step 2. diff --git a/src/rust/cxx/kj-rs/tests/test_own.rs b/src/rust/cxx/kj-rs/tests/test_own.rs index 61d0235b08c..292c3950ed7 100644 --- a/src/rust/cxx/kj-rs/tests/test_own.rs +++ b/src/rust/cxx/kj-rs/tests/test_own.rs @@ -201,20 +201,10 @@ pub mod tests { assert!(!debug_str.is_empty()); } - #[test] - fn test_own_send_between_threads() { - use std::thread; - - let own = ffi::cxx_kj_own(); - let handle = thread::spawn(move || { - // Own should be Send, so this should work - assert_eq!(own.get_data(), 42); - own - }); - - let returned_own = handle.join().unwrap(); - assert_eq!(returned_own.get_data(), 42); - } + // NOTE: `KjOwn` is deliberately neither `Send` nor `Sync`, even for `T: Send + Sync` + // (like `OpaqueCxxClass` here): the `KjOwn` carries a type-erased C++ disposer which may + // not be thread-safe (kj::Rc, arena allocations, ...). The negative is asserted at compile + // time in lib.rs (assert_not_impl_any). #[test] fn test_own_concurrent_creation() { @@ -254,49 +244,4 @@ pub mod tests { let expected: Vec = (0..num_threads).map(|i| i as u64 * 100).collect(); assert_eq!(results, expected); } - - // This is one test generated by Claude. I am unsure it sufficiently tests multithreading. - #[test] - fn test_own_stress_multithreaded() { - use std::sync::mpsc; - use std::thread; - - let (tx, rx) = mpsc::channel(); - let num_threads: u64 = 12; - let items_per_thread: u64 = 100; - - for thread_id in 0..num_threads { - let tx_clone = tx.clone(); - thread::spawn(move || { - for i in 0..items_per_thread { - let mut own = ffi::cxx_kj_own(); - let value = thread_id * items_per_thread + i; - own.pin_mut().set_data(value); - - // Send the Own across thread boundary - tx_clone.send(own).unwrap(); - } - }); - } - drop(tx); // Close the sending side - - // Collect all Owns from all threads - let mut received_owns = Vec::new(); - while let Ok(own) = rx.recv() { - received_owns.push(own); - } - - // Verify we received the expected number - assert_eq!( - received_owns.len(), - (num_threads * items_per_thread) as usize - ); - - // Verify all values are correct - let mut values: Vec = received_owns.iter().map(|own| own.get_data()).collect(); - values.sort_unstable(); - - let expected: Vec = (0..(num_threads * items_per_thread)).collect(); - assert_eq!(values, expected); - } } diff --git a/src/rust/cxx/kj-rs/waker.c++ b/src/rust/cxx/kj-rs/waker.c++ index 0c993fdaf7e..b0311325db7 100644 --- a/src/rust/cxx/kj-rs/waker.c++ +++ b/src/rust/cxx/kj-rs/waker.c++ @@ -1,150 +1,48 @@ #include "waker.h" -#include +#include "awaiter.h" namespace kj_rs { -// ======================================================================================= -// ArcWakerPromiseNode - -ArcWakerPromiseNode::ArcWakerPromiseNode(kj::Promise promise) - : node(PromiseNode::from(kj::mv(promise))) { - node->setSelfPointer(&node); -} - -void ArcWakerPromiseNode::destroy() noexcept { - auto drop = kj::mv(owner); -} - -void ArcWakerPromiseNode::onReady(kj::_::Event* event) noexcept { - node->onReady(event); -} - -void ArcWakerPromiseNode::get(kj::_::ExceptionOrValue& output) noexcept { - node->get(output); - KJ_IF_SOME(exception, kj::runCatchingExceptions([this]() { node = nullptr; })) { - output.addException(kj::mv(exception)); - } -} - -void ArcWakerPromiseNode::tracePromise(kj::_::TraceBuilder& builder, bool stopAtNextEvent) { - // TODO(someday): Is it possible to get the address of the Rust code which cloned our Waker? - - if (node.get() != nullptr) { - node->tracePromise(builder, stopAtNextEvent); - } -} - -// ======================================================================================= -// ArcWaker - -PromiseArcWakerPair ArcWaker::create(const kj::Executor& executor) { - // TODO(perf): newPromiseAndCrossThreadFulfiller() makes two heap allocations, but it is probably - // optimizable to one. - // TODO(perf): This heap allocation could also probably be collapsed into the fulfiller's. - auto waker = - kj::arc(kj::Badge(), executor.newPromiseAndCrossThreadFulfiller()); - auto promise = const_cast(waker.get())->getPromise(); - return { - .promise = kj::mv(promise), - .waker = kj::mv(waker), - }; -} - -kj::Promise ArcWaker::getPromise() { - KJ_REQUIRE(node.owner == nullptr); - node.owner = addRefToThis(); - return kj::_::PromiseNode::to>(OwnPromiseNode(&node)); -} - -ArcWaker::ArcWaker(kj::Badge, kj::PromiseCrossThreadFulfillerPair paf) - : node(kj::mv(paf.promise)), - fulfiller(kj::mv(paf.fulfiller)) {} - -const KjWaker* ArcWaker::clone() const { - return addRefToThis().disown(); -} -void ArcWaker::wake() const { - wake_by_ref(); - drop(); -} -void ArcWaker::wake_by_ref() const { - fulfiller->fulfill(); -} -void ArcWaker::drop() const { - auto drop = kj::Arc::reown(this); -} +// Definition of the arm-nudge hook declared in waker.h; null until an integrating event port +// (kj-rs-tokio's TokioEventPort) installs itself. Thread-local: one loop/port per thread. +thread_local void (*futurePollArmNudge)() = nullptr; // ======================================================================================= -// LazyArcWaker - -const KjWaker* LazyArcWaker::clone() const { - // Rust code wants to suspend and wait for something. We'll start handing out ArcWakers if we - // haven't already been woken synchronously. - - if (wakeCount.load(std::memory_order_relaxed) > 0) { - // We were already woken synchronously, so there's no point handing out more wakers for the - // current call to `Future::poll()`. We can hand out a noop waker by returning nullptr. - return nullptr; - } - - auto lock = cloned.lockExclusive(); - - if (*lock == kj::none) { - // We haven't been cloned before, so make a new ArcWaker. - *lock = ArcWaker::create(executor); +// PollWaker +// +// These are defined here rather than inline in waker.h because they reach into FuturePollEvent, +// which is only a complete type once awaiter.h is included. + +PollWaker::PollWaker(FuturePollEvent& futurePollEvent) + : holder(FuturePollEventHolder{futurePollEvent}) {} + +PollWaker::~PollWaker() noexcept(false) {} + +void PollWaker::wakeByRef() const { + // Synchronous same-turn wake during `future.poll()`: arm the FuturePollEvent so it re-polls. + // armDepthFirst() is idempotent and safe to call from within the event's own fire(), so this + // works whether we were reached from onReady() (first poll) or fire() (a subsequent poll). No + // arm-nudge needed here: the loop is running this very poll, not parked. + KJ_IF_SOME(futurePollEvent, tryGetFuturePollEvent()) { + futurePollEvent.armDepthFirst(); } - - return KJ_ASSERT_NONNULL(*lock).waker->clone(); -} - -void LazyArcWaker::wake() const { - // LazyArcWakers are only exposed to Rust by const borrow, meaning Rust can never arrange to call - // `wake()`, which drops `self`, on this object. - KJ_UNIMPLEMENTED("Rust user code should never have possess a consumable " - "reference to LazyArcWaker"); -} - -void LazyArcWaker::wake_by_ref() const { - // Woken synchronously during a call to `future.poll(awaitWaker)`. - wakeCount.fetch_add(1, std::memory_order_relaxed); } -void LazyArcWaker::drop() const { - ++dropCount; +kj::Maybe> PollWaker::cloneCell() const { + // Rust wants a waker it can retain and wake later. Hand out a strong reference to a + // FutureWakerCell bound to the FuturePollEvent being polled; waking it arms that event. + KJ_IF_SOME(futurePollEvent, tryGetFuturePollEvent()) { + return futurePollEvent.cloneWakerCell(); + } + return kj::none; } -kj::Maybe> LazyArcWaker::reset() { - // This function is only called after `future.poll(awaitWaker)` has returned, meaning Rust has - // dropped its reference. Thus, we don't need to worry about thread-safety here, and can call - // `cloned.getWithoutLock()`, for example. - - KJ_ASSERT(dropCount == 1); - KJ_DEFER(dropCount = 0); - KJ_DEFER(wakeCount.store(0, std::memory_order_relaxed)); - - // Reset the ArcWaker on our way out. Since we only return the ArcWaker's promise to our caller, - // we ensure that Rust owns the only remaining ArcWaker clones, if any. - // - // TODO(perf): If ArcWakers were resettable, we could instead return the ArcWaker for our caller - // to cache for later use. - KJ_DEFER(cloned.getWithoutLock() = kj::none); - - if (wakeCount.load(std::memory_order_relaxed) > 0) { - // The future returned Pending, but synchronously called `wake_by_ref()` on the LazyArcWaker, - // indicating it wants to immediately be polled again. We should arm our event right now, - // which will call `await_ready()` again on the event loop. - return kj::Promise(kj::READY_NOW); - } else KJ_IF_SOME(arcWakerPair, cloned.getWithoutLock()) { - // The future returned Pending and cloned an ArcWaker to notify us later. We'll arrange for - // the ArcWaker's promise to arm our event once it's fulfilled. - return kj::mv(arcWakerPair.promise); - } else { - // The future returned Pending, did not call `wake_by_ref()` on the LazyArcWaker, and did not - // clone an ArcWaker. Rust is either awaiting a KJ promise, or the Rust equivalent of - // kj::NEVER_DONE. - return kj::none; +kj::Maybe PollWaker::tryGetFuturePollEvent() const { + KJ_IF_SOME(h, holder.tryGet()) { + return h.futurePollEvent; } + return kj::none; } } // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/waker.h b/src/rust/cxx/kj-rs/waker.h index 382572aabb4..92cf54921fe 100644 --- a/src/rust/cxx/kj-rs/waker.h +++ b/src/rust/cxx/kj-rs/waker.h @@ -1,176 +1,132 @@ #pragma once -#include "promise.h" +#include "kj-rs/executor-guarded.h" #include -#include -#include #include -#include - namespace kj_rs { -using kj::uint; +class FuturePollEvent; -// ======================================================================================= -// KjWaker +// Hook invoked when a waker arms its FuturePollEvent (below). An integrating kj::EventPort that +// drives tokio tasks inside its own wait() (see kj-rs-tokio) installs this to nudge itself out of +// a blocking park: a tokio task completing during the port's block_on() may arm a KJ event +// same-thread, and KJ's edge-triggered setRunnable() misses that arm when the loop's runnable +// state is already set (e.g. left true by a prior timer). Null (no-op) by default; thread-local, +// one integrating port per loop thread. +extern thread_local void (*futurePollArmNudge)(); -class FuturePollEvent; +// ======================================================================================= +// FutureWakerCell -// KjWaker is an abstract base class which defines an interface mirroring Rust's RawWakerVTable -// struct. Rust has four trampoline functions, defined in waker.rs, which translate Waker::clone(), -// Waker::wake(), etc. calls to the virtual member functions on this class. +// FutureWakerCell is the same-thread "waker cell" behind every Rust waker that outlives a single +// `Future::poll()` call, refcounted via the non-atomic kj::Refcounted (single-thread axiom: a +// waker is never woken, cloned, or dropped from another thread). +// +// The cell holds a bare pointer to the owning FuturePollEvent's `kj::_::Event`. `wakeByRef()` +// arms that Event directly via `Event::armDepthFirst()`, which is idempotent across same-turn +// arms and safe to call from within the event's own `fire()` (turn() unlinks the event before +// firing). This is the same way a RustPromiseAwaiter leaf arms the FuturePollEvent when its +// Promise becomes ready — a wake is just another leaf arming the same event. +// +// Neutralize-on-drop: the cell's link to the Event is weak — a `kj::Maybe` invalidated +// (structurally, by the owning FuturePollEvent's RAII guard; see awaiter.h) when that event is +// destroyed. A cell reference that Rust retains past the Future's lifetime (e.g. parked in a +// channel's AtomicWaker) therefore observes a dead link on a later wake and is a safe no-op, +// rather than arming a freed Event. // -// Rust requires Wakers to be Send and Sync, meaning all of the functions defined here may be called -// concurrently by any thread. Derived class implementations of these functions must handle this, -// which is why all of the virtual member functions are `const`-qualified. -class KjWaker { +// Ownership only ever crosses the FFI as real `kj::Rc` handles (PollWaker:: +// cloneCell(), addRef()). Rust's RawWakerVTable island (waker.rs) carries its handle in the +// RawWaker data slot, disowning/reowning it at the vtable edge — the one place `std::task::Waker` +// forces a raw pointer. +class FutureWakerCell final: public kj::Refcounted { public: - // Return a pointer to a new strong ref to a KjWaker. Note that `clone()` may return nullptr, - // in which case the Rust implementation in waker.rs will treat it as a no-op Waker. Rust - // immediately wraps this pointer in its own Waker object, which is responsible for later - // releasing the strong reference. - // - // TODO(cleanup): Build kj::Arc into cxx-rs so we can return one instead of a raw pointer. - virtual const KjWaker* clone() const = 0; - - // Wake and drop this waker. - virtual void wake() const = 0; - - // Wake this waker, but do not drop it. - virtual void wake_by_ref() const = 0; - - // Drop this waker. - virtual void drop() const = 0; - - // If this KjWaker implementation has an associated FuturePollEvent, C++ code can request access - // to it here. The RustPromiseAwaiter class (which helps Rust `.await` KJ Promises) uses this to - // optimize awaits, when possible. - virtual kj::Maybe tryGetFuturePollEvent() const { - return kj::none; - } -}; + explicit FutureWakerCell(kj::_::Event& event): event(event) {} -// ======================================================================================= -// ArcWakerPromiseNode + // Called by `~FuturePollEvent` to neutralize this cell and every outstanding Rust reference to + // it, making any subsequent wake a safe no-op. + void neutralize() { + event = kj::none; + } -class ArcWaker; + // Arm the owning FuturePollEvent, or no-op if it has been neutralized. Const because Rust + // reaches it through `&self`; in the single-thread world it only ever runs on the owning event + // loop's thread. + void wakeByRef() const { + KJ_IF_SOME(e, event) { + e.armDepthFirst(); + // Nudge an integrating event port out of a blocking park (see `futurePollArmNudge`). No-op + // unless a port installed the hook and is currently parked. + if (futurePollArmNudge != nullptr) { + futurePollArmNudge(); + } + } + } -class ArcWakerPromiseNode: public kj::_::PromiseNode { - public: - ArcWakerPromiseNode(kj::Promise promise); - KJ_DISALLOW_COPY_AND_MOVE(ArcWakerPromiseNode); + // Hand out a new strong reference. Const + const_cast because Rust reaches it through `&self`: + // cells are always heap-allocated non-const (kj::rc in cloneWakerCell()), and the non-atomic + // refcount bump is safe under the single-thread axiom. + kj::Rc addRef() const { + return const_cast(*this).addRefToThis(); + } - void destroy() noexcept override; - void onReady(kj::_::Event* event) noexcept override; - void get(kj::_::ExceptionOrValue& output) noexcept override; - void tracePromise(kj::_::TraceBuilder& builder, bool stopAtNextEvent) override; + // Re-own a strong reference previously surrendered to a raw pointer: waker.rs disowns the + // kj::Rc it parks in a RawWaker data slot, and its vtable's drop calls this to reclaim it. + // Exposed to Rust as an `unsafe fn`: `this` must carry exactly such a surrendered reference, + // and dropping the returned handle releases it. + kj::Rc reown() const { + return kj::Rc::reown(&const_cast(*this)); + } private: - kj::Arc owner = nullptr; - OwnPromiseNode node; - - friend class ArcWaker; + // Weak, owner-invalidated reference to the owning FuturePollEvent's Event base: non-owning (the + // event lives in the promise graph; the cell must observe its death, never extend its life) and + // nulled by `neutralize()` when that event is destroyed. Reads in the `const` wake functions + // are single-thread so no synchronization is required. + kj::Maybe event; }; // ======================================================================================= -// ArcWaker - -class ArcWaker; +// PollWaker -struct PromiseArcWakerPair { - kj::Promise promise; - kj::Arc waker; -}; - -// ArcWaker is an atomic-refcounted wrapper around a `CrossThreadPromiseFulfiller`. -// The atomic-refcounted aspect makes it safe to call `clone()` and `drop()` concurrently, while the -// `CrossThreadPromiseFulfiller` aspect makes it safe to call `wake_by_ref()` concurrently. Finally, -// `wake()` is implemented in terms of `wake_by_ref()` and `drop()`. +// PollWaker is the waker C++ passes to `Future::poll()`. It lives on the stack / in a coroutine +// frame for the duration of a single poll, and Rust only ever borrows it (waker.rs wraps it in a +// Waker whose drop is a no-op). // -// This class is mostly an implementation detail of LazyArcWaker. -class ArcWaker: public kj::AtomicRefcounted, public KjWaker { +// - wakeByRef() is a synchronous same-turn wake: it arms the FuturePollEvent directly. Arming +// during poll() (whether poll was reached via onReady() or fire()) is idempotent and causes +// an immediate re-poll. +// - cloneCell() is how Rust retains a waker past the poll: it hands out a strong reference to +// the event's FutureWakerCell, so a later wake arms the same event. +// - tryGetFuturePollEvent() lets RustPromiseAwaiter (which helps Rust `.await` KJ Promises) +// arm the event directly instead of going through a waker, when possible. +class PollWaker final { public: - // Construct a new promise and ArcWaker promise pair, with the Promise to be scheduled on the - // event loop associated with `executor`. - static PromiseArcWakerPair create(const kj::Executor& executor); + // `futurePollEvent` is the FuturePollEvent responsible for calling `Future::poll()`, and must + // outlive this PollWaker. + explicit PollWaker(FuturePollEvent& futurePollEvent); + ~PollWaker() noexcept(false); + KJ_DISALLOW_COPY_AND_MOVE(PollWaker); - ArcWaker(kj::Badge, kj::PromiseCrossThreadFulfillerPair paf); - KJ_DISALLOW_COPY_AND_MOVE(ArcWaker); + // Synchronous same-turn wake: arm the associated FuturePollEvent so it re-polls. + void wakeByRef() const; - const KjWaker* clone() const override; - void wake() const override; - void wake_by_ref() const override; - void drop() const override; + // Get-or-create the event's FutureWakerCell and hand out a new strong reference to it, for Rust + // to retain and wake later. Returns kj::none if the current thread's kj::Executor is not the + // one which owns the FuturePollEvent (cannot normally happen in the single-thread world); Rust + // then mints a no-op waker. + kj::Maybe> cloneCell() const; - private: - kj::Promise getPromise(); - - ArcWakerPromiseNode node; - kj::Own> fulfiller; -}; - -// ======================================================================================= -// LazyArcWaker - -// LazyArcWaker is intended to live locally on the stack or in a coroutine frame. Trying to -// `clone()` it will cause it to allocate an ArcWaker for the caller. -class LazyArcWaker: public KjWaker { - public: - // Create a new or clone an existing ArcWaker, leak its pointer, and return it. This may be called - // by any thread. - const KjWaker* clone() const override; - - // Unimplemented, because Rust user code cannot consume the `std::task::Waker` we create which - // wraps this LazyArcWaker. - void wake() const override; - - // Rust user code can wake us synchronously during the execution of `future.poll()` using this - // function. This may be called by any thread. - void wake_by_ref() const override; - - // Does not actually destroy this object. Instead, we increment a counter so we can assert that it - // was dropped exactly once before `future.poll()` returned. This can only be called on the thread - // which is doing the awaiting, because our implementation of `future.poll()` never transfers the - // Waker object to a different thread. - void drop() const override; - - // Used by the owner of LazyArcWaker after `future.poll()` has returned, to retrieve the - // LazyArcWaker's state for further processing. This is non-const, because by the time this is - // called, Rust has dropped all of its borrows to this class, meaning we no longer have to worry - // about thread safety. - // - // This function will assert if `drop()` has not been called since LazyArcWaker was constructed, - // or since the last call to `reset()`. - // - // Returns `kj::none` the LazyArcWaker was neither woken nor cloned before being dropped. Returns - // `kj::READY_NOW` if the LazyArcWaker was synchronously woken. Otherwise, if `clone()` was - // called, return the promise associated with the cloned ArcWaker. - kj::Maybe> reset(); + // The FuturePollEvent whose poll() this waker was created for, if the current thread's + // kj::Executor is the one which owns it. + kj::Maybe tryGetFuturePollEvent() const; private: - // We store the kj::Executor for the constructing thread so that we can lazily instantiate a - // CrossThreadPromiseFulfiller from any thread in our `clone()` implementation. - const kj::Executor& executor = kj::getCurrentThreadExecutor(); - - // Initialized by `clone()`, which may be called by any thread. This could almost be a - // `kj::Lazy`, but we need to be able to detect when we haven't been cloned. - kj::MutexGuarded> cloned; - - // Incremented by `wake_by_ref()`, which may be called by any thread. All operations use relaxed - // memory order, because this counter does not guard any memory. - mutable std::atomic wakeCount{0}; - - // Incremented by `drop()`, so we can validate that `drop()` is only called once on this object. - // - // Rust requires that Wakers be droppable by any thread. However, we own the implementation of - // `poll()` to which `LazyArcWaker&` is passed, and those implementations store the Rust - // `std::task::Waker` object on the stack,, and never move it elsewhere. Since that object is - // responsible for calling `LazyArcWaker::drop()`, we know for sure that `drop()` will only ever be - // called on the thread which constructed it. Therefore, there is no need to make `dropCount` - // thread-safe. - mutable uint dropCount = 0; + struct FuturePollEventHolder { + FuturePollEvent& futurePollEvent; + }; + ExecutorGuarded holder; }; } // namespace kj_rs diff --git a/src/rust/cxx/kj-rs/waker.rs b/src/rust/cxx/kj-rs/waker.rs index 472a259aaa3..7fdf482a53a 100644 --- a/src/rust/cxx/kj-rs/waker.rs +++ b/src/rust/cxx/kj-rs/waker.rs @@ -1,102 +1,184 @@ +//! FFI island (see crate-root `#![deny(unsafe_code)]`): the two `RawWakerVTable`s bridging the +//! C++ wakers into `std::task::Waker`. `std`'s vtable ABI is four raw-pointer functions, so this +//! file is where ownership must round-trip through a raw `RawWaker` data slot — everywhere else +//! (including the whole C++ interface) waker ownership is a real `kj::Rc` handle. A genuine unsafe +//! seam. +#![allow(unsafe_code)] + use std::task::RawWaker; use std::task::RawWakerVTable; use std::task::Waker; -use crate::ffi::KjWaker; +use crate::KjRc; +use crate::ffi::FutureWakerCell; +use crate::ffi::PollWaker; -// Safety: We use the type system to express the Sync nature of KjWaker in the cxx-rs FFI boundary. -// Specifically, we only allow invocations on const KjWakers, and in KJ C++, use of const-qualified -// functions is thread-safe by convention. Our implementations of KjWakers in C++ respect this -// convention. -// -// Note: Implementing these traits does not seem to be required for building, but the Waker -// documentation makes it clear Send and Sync are a requirement of the pointed-to type. +// Thread-safety: `std::task::Waker` documents that the vtable functions must be thread-safe. +// The bridge is single-threaded — no waker is ever woken, cloned, or dropped from another thread +// (cross-thread producers interpose a same-thread forwarding task instead; see kj-rs-io's +// `resolve_host`) — so these vtables uphold that contract degenerately: every call happens on the +// owning event loop's thread. + +// ======================================================================================= +// Borrowed vtable: Wakers lending out the PollWaker C++ passes to `Future::poll()` // -// https://doc.rust-lang.org/std/task/struct.RawWaker.html -// https://doc.rust-lang.org/std/task/struct.RawWakerVTable.html -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. -unsafe impl Send for KjWaker {} -// Safety: the KJ bridge representation and ownership invariants satisfy this operation. -unsafe impl Sync for KjWaker {} - -impl From<&KjWaker> for Waker { - fn from(waker: &KjWaker) -> Self { - let waker = RawWaker::new( - std::ptr::from_ref::(waker).cast::<()>(), - &KJ_WAKER_VTABLE, +// `data` is the `&PollWaker` the Waker was built from — borrowed, never null, and alive for the +// duration of the poll (the Waker is created and dropped inside the poll bridge in future.rs). +// Dropping such a Waker frees nothing; cloning it takes a real strong reference to the event's +// FutureWakerCell and switches to the owned-cell vtable below. + +impl From<&PollWaker> for Waker { + fn from(waker: &PollWaker) -> Self { + let raw = RawWaker::new( + std::ptr::from_ref::(waker).cast::<()>(), + &POLL_WAKER_VTABLE, ); - // Safety: KjWaker's Rust-exposed interface is Send and Sync and its RawWakerVTable - // implementation functions are all thread-safe. - // - // https://doc.rust-lang.org/std/task/struct.Waker.html#safety-1 - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - unsafe { Self::from_raw(waker) } + // Safety: the vtable functions below uphold the RawWaker contract (see the thread-safety + // note above); `data` outlives the Waker because future.rs drops the Waker before poll + // returns. + unsafe { Self::from_raw(raw) } } } -// Helper function for use in KjWaker's RawWakerVTable implementation to factor out a tedious null -// pointer check. -fn deref_kj_waker<'a>(data: *const ()) -> Option<&'a KjWaker> { +/// # Safety +/// +/// `data` must be the pointer a [`From<&PollWaker>`] conversion was made with, still live per the +/// `RawWaker` contract (upheld because these Wakers only exist within a single `poll` call). +unsafe fn poll_waker_clone(data: *const ()) -> RawWaker { + // Safety: forwarded from this fn's `# Safety` contract. + let waker = unsafe { &*data.cast::() }; + let cell: Option> = waker.clone_cell().into(); + RawWaker::new( + cell.map_or(std::ptr::null(), cell_into_raw).cast::<()>(), + &CELL_WAKER_VTABLE, + ) +} + +/// # Safety +/// +/// Same contract as [`poll_waker_clone`]. +unsafe fn poll_waker_wake_by_ref(data: *const ()) { + // Safety: forwarded from this fn's `# Safety` contract. + let waker = unsafe { &*data.cast::() }; + waker.wake_by_ref(); +} + +/// # Safety +/// +/// Same contract as [`poll_waker_clone`]. Consuming a borrowed Waker owns nothing, so `wake` is +/// just `wake_by_ref` (the paired drop is a no-op). +unsafe fn poll_waker_wake(data: *const ()) { + // Safety: forwarded from this fn's `# Safety` contract. + unsafe { poll_waker_wake_by_ref(data) } +} + +fn poll_waker_drop(_data: *const ()) { + // No-op: the PollWaker is stack-owned by the C++ poll scope; this Waker only borrowed it. +} + +static POLL_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new( + poll_waker_clone, + poll_waker_wake, + poll_waker_wake_by_ref, + poll_waker_drop, +); + +// ======================================================================================= +// Owned-cell vtable: retained Wakers holding a strong reference to a FutureWakerCell +// +// `data` carries one strong reference (a disowned `KjRc`), or is null for the +// no-op Waker minted when `clone_cell()` had no event to bind (cannot normally happen in the +// single-thread world). Clone takes another reference; drop re-owns and releases the carried one. + +/// Surrender the handle's strong reference into a bare pointer for a `RawWaker` data slot. +/// Reversed by `FutureWakerCell::reown` in [`cell_waker_drop`]. +fn cell_into_raw(cell: KjRc) -> *const FutureWakerCell { + let ptr = cell.get(); + std::mem::forget(cell); + ptr +} + +/// # Safety +/// +/// `data` must be either null or a pointer produced by [`cell_into_raw`] whose strong reference +/// is still carried by this `RawWaker` (upheld by the `Waker`/`RawWaker` contract: these vtable +/// entries are only installed alongside such pointers, by [`poll_waker_clone`] and the functions +/// below). +unsafe fn cell_deref<'a>(data: *const ()) -> Option<&'a FutureWakerCell> { if data.is_null() { None } else { - let p = data.cast::(); - // Safety: - // 1. p is guaranteed non-null by the check above. - // 2. This function is only used in the implementations of our RawWakerVTable for KjWaker. - // All vtable implementation functions are trivially guaranteed that their owning Waker - // object is still alive. We assume the Waker was constructed correctly to begin with, - // and that therefore the pointer still points to valid memory. - // 3. We do not read or write the KjWaker's memory, so there are no atomicity concerns nor - // interleaved pointer/reference access concerns. - // - // https://doc.rust-lang.org/std/ptr/index.html#safety - // Safety: the KJ bridge representation and ownership invariants satisfy this operation. - Some(unsafe { &*p }) + // Safety: non-null per the check; live per this fn's `# Safety` contract (the carried + // strong reference keeps the cell alive). + Some(unsafe { &*data.cast::() }) } } -pub fn kj_waker_clone(data: *const ()) -> RawWaker { - let new_data = if let Some(kj_waker) = deref_kj_waker(data) { - kj_waker.clone_kj_waker().cast::<()>() +/// # Safety +/// +/// Same contract as [`cell_deref`]. +unsafe fn cell_waker_clone(data: *const ()) -> RawWaker { + // Safety: forwarded from this fn's `# Safety` contract. + let new_data = if let Some(cell) = unsafe { cell_deref(data) } { + cell_into_raw(cell.add_ref()) } else { std::ptr::null() }; - RawWaker::new(new_data, &KJ_WAKER_VTABLE) + RawWaker::new(new_data.cast::<()>(), &CELL_WAKER_VTABLE) } -pub fn kj_waker_wake(data: *const ()) { - if let Some(kj_waker) = deref_kj_waker(data) { - kj_waker.wake(); +/// # Safety +/// +/// Same contract as [`cell_deref`]. +unsafe fn cell_waker_wake_by_ref(data: *const ()) { + // Safety: forwarded from this fn's `# Safety` contract. + if let Some(cell) = unsafe { cell_deref(data) } { + cell.wake_by_ref(); } } -pub fn kj_waker_wake_by_ref(data: *const ()) { - if let Some(kj_waker) = deref_kj_waker(data) { - kj_waker.wake_by_ref(); +/// # Safety +/// +/// Same contract as [`cell_deref`], and the carried strong reference is released (the `RawWaker` +/// must not be used again — guaranteed by the `Waker` contract for `drop`). +unsafe fn cell_waker_drop(data: *const ()) { + // Safety: forwarded from this fn's `# Safety` contract. + if let Some(cell) = unsafe { cell_deref(data) } { + // Safety: `data` carries a strong reference per this fn's `# Safety` contract; re-own it + // and let the handle fall, releasing the reference. + let _cell = unsafe { cell.reown() }; } } -pub fn kj_waker_drop(data: *const ()) { - if let Some(kj_waker) = deref_kj_waker(data) { - kj_waker.drop(); +/// # Safety +/// +/// Same contract as [`cell_waker_drop`]. +unsafe fn cell_waker_wake(data: *const ()) { + // Safety: forwarded from this fn's `# Safety` contract; wake-then-release. + unsafe { + cell_waker_wake_by_ref(data); + cell_waker_drop(data); } } -static KJ_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new( - kj_waker_clone, - kj_waker_wake, - kj_waker_wake_by_ref, - kj_waker_drop, +static CELL_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new( + cell_waker_clone, + cell_waker_wake, + cell_waker_wake_by_ref, + cell_waker_drop, ); -/// If `waker` wraps a `KjWaker`, return the `KjWaker` pointer it was originally constructed with, -/// or null if `waker` does not wrap a `KjWaker`. Note that the `KjWaker` pointer originally used -/// to construct `waker` may itself by null. -pub fn try_into_kj_waker_ptr(waker: &Waker) -> *const KjWaker { - if waker.vtable() == &KJ_WAKER_VTABLE { - waker.data().cast::() +/// If `waker` lends out a C++ `PollWaker` (borrowed vtable above), return a reference to it, +/// borrowed from `waker` itself. Owned-cell and foreign Wakers both return `None`: neither +/// exposes a `FuturePollEvent` to arm directly, so `RustPromiseAwaiter` takes its generic +/// fallback path for them. +pub fn try_poll_waker(waker: &Waker) -> Option<&PollWaker> { + if waker.vtable() == &POLL_WAKER_VTABLE { + // Safety: Wakers carrying POLL_WAKER_VTABLE are only ever built by `From<&PollWaker>` + // above, so `data` is a `&PollWaker` that outlives `waker` (the PollWaker is stack-owned + // by the C++ poll driving this call); the returned borrow is tied to `waker`'s lifetime. + Some(unsafe { &*waker.data().cast::() }) } else { - std::ptr::null() + None } } diff --git a/src/rust/kj/tests/ffi-test.c++ b/src/rust/kj/tests/ffi-test.c++ index e4218071c45..34cee77d7ee 100644 --- a/src/rust/kj/tests/ffi-test.c++ +++ b/src/rust/kj/tests/ffi-test.c++ @@ -48,10 +48,11 @@ class MockHttpService: public kj::HttpService { class TestConnectResponse: public kj::HttpService::ConnectResponse { public: - void accept(uint statusCode, kj::StringPtr statusText, const kj::HttpHeaders& headers) override { + void accept( + kj::uint statusCode, kj::StringPtr statusText, const kj::HttpHeaders& headers) override { KJ_UNIMPLEMENTED("not exercised by test"); } - kj::Own reject(uint statusCode, + kj::Own reject(kj::uint statusCode, kj::StringPtr statusText, const kj::HttpHeaders& headers, kj::Maybe expectedBodySize = kj::none) override { diff --git a/src/workerd/api/BUILD.bazel b/src/workerd/api/BUILD.bazel index 76b5f37c6ba..d1eb7bcdead 100644 --- a/src/workerd/api/BUILD.bazel +++ b/src/workerd/api/BUILD.bazel @@ -417,7 +417,7 @@ wd_cc_library( "//src/pyodide:python_packages_capnp", "//src/workerd/io:compatibility-date_capnp", "//src/workerd/jsg", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", "@capnp-cpp//src/kj/compat:kj-gzip", "@capnp-cpp//src/kj/compat:kj-http", "@ssl", @@ -495,7 +495,7 @@ wd_cc_library( visibility = ["//visibility:public"], deps = [ "@capnp-cpp//src/kj", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", ], ) @@ -767,6 +767,7 @@ kj_test( "//src/workerd/io", "//src/workerd/jsg", "//src/workerd/tests:test-fixture", + "//src/workerd/util:setup-async-io", ], ) diff --git a/src/workerd/io/BUILD.bazel b/src/workerd/io/BUILD.bazel index ee64c29c3d7..be101412930 100644 --- a/src/workerd/io/BUILD.bazel +++ b/src/workerd/io/BUILD.bazel @@ -124,7 +124,7 @@ wd_cc_library( "//src/workerd/util:strong-bool", "@capnp-cpp//src/capnp:capnp-rpc", "@capnp-cpp//src/capnp/compat:http-over-capnp", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", "@ncrypto", "@ssl", ], @@ -156,7 +156,7 @@ wd_cc_library( ":worker-interface_capnp", "//src/workerd/jsg", "@capnp-cpp//src/capnp/compat:http-over-capnp", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", ], ) @@ -247,7 +247,7 @@ wd_cc_library( "//src/workerd/util:sqlite", "//src/workerd/util:strong-bool", "@capnp-cpp//src/capnp:capnp-rpc", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", ], ) @@ -313,7 +313,7 @@ wd_cc_library( deps = [ ":trace", "@capnp-cpp//src/kj", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", ], ) @@ -325,7 +325,7 @@ wd_cc_library( visibility = ["//visibility:public"], deps = [ "@capnp-cpp//src/kj", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", ], ) @@ -619,6 +619,7 @@ kj_test( deps = [ ":io", "//src/workerd/tests:test-fixture", + "//src/workerd/util:setup-async-io", ], ) diff --git a/src/workerd/jsg/BUILD.bazel b/src/workerd/jsg/BUILD.bazel index 9d9eb74f984..27b960d1ae3 100644 --- a/src/workerd/jsg/BUILD.bazel +++ b/src/workerd/jsg/BUILD.bazel @@ -313,6 +313,7 @@ wd_cc_library( exclude = [ # defined below "macro-meta-test.c++", + "modules-new-test.c++", "resource-test.c++", "rtti-test.c++", "url-test.c++", @@ -320,6 +321,19 @@ wd_cc_library( ], )] +# Pulled out of the glob above because it calls kj::setupAsyncIo(); :jsg only pulls the core kj +# target, so it depends on the io-backend seam (//src/workerd/util:setup-async-io) which supplies +# kj::setupAsyncIo() for whichever backend is selected -- native under cxx, tokio under +# --//:io_backend=rust -- with no #if in the test source. +kj_test( + src = "modules-new-test.c++", + local_defines = ["JSG_IMPLEMENTATION"], + deps = [ + ":jsg", + "//src/workerd/util:setup-async-io", + ], +) + # Moved out as macro-meta-test does not depend on V8 or JSG proper, this makes the test much # smaller. kj_test( diff --git a/src/workerd/server/BUILD.bazel b/src/workerd/server/BUILD.bazel index 22284af7725..7d70e8eff28 100644 --- a/src/workerd/server/BUILD.bazel +++ b/src/workerd/server/BUILD.bazel @@ -2,6 +2,7 @@ load("@bazel_skylib//lib:selects.bzl", "selects") load("@bazel_skylib//rules:common_settings.bzl", "bool_flag") load("@bazel_skylib//rules:copy_file.bzl", "copy_file") load("//:build/kj_test.bzl", "kj_test") +load("//:build/rust_io_backend.bzl", "rust_io_backend_local_defines", "rust_io_hermeticity") load("//:build/wd_capnp_library.bzl", "wd_capnp_library") load("//:build/wd_cc_binary.bzl", "wd_cc_binary") load("//:build/wd_cc_embed.bzl", "wd_cc_embed") @@ -67,6 +68,7 @@ wd_cc_binary( malloc = ":malloc", visibility = ["//visibility:public"], deps = [ + ":cli-io-backend", ":cpp-capnp-schema", ":json-logger", ":server", @@ -74,13 +76,54 @@ wd_cc_binary( ":workerd-capnp-schema", ":workerd_capnp", "//src/pyodide:pyodide_extra_capnp", - "//src/rust/cxx-integration", "//src/workerd/util:autogate", "//src/workerd/util:perfetto", + "//src/workerd/util:setup-async-io", "@capnp-cpp//src/capnp:capnpc", ], ) +# The CLI main()'s two backend-divergent I/O paths -- the --watch file watcher and the SIGTERM +# graceful drain -- behind backend-agnostic entry points (workerd::server::makeFileWatcher / +# captureSigterm / onSigterm), so workerd.c++ links them unchanged in both configs. The single +# WORKERD_RUST_IO_BACKEND_RUST #if selecting the native kj loop vs the tokio loop lives entirely in +# cli-io-backend.c++, mirroring //src/workerd/util:setup-async-io. Under --//:io_backend=rust this +# supplies the tokio (kj-rs-io) backends; in the default cxx build it uses the native +# kj::UnixEventPort watcher + signal handling. +wd_cc_library( + name = "cli-io-backend", + srcs = ["cli-io-backend.c++"], + hdrs = ["cli-io-backend.h"], + local_defines = rust_io_backend_local_defines(), + visibility = ["//visibility:public"], + deps = select({ + "//:io_backend_rust": [ + "//deps:rust_runtime", + "//src/rust/cxx/kj-rs-io", + "//src/rust/cxx/kj-rs-io:kj-rs-io-lib", + "@capnp-cpp//src/kj:kj-async-core", + ], + "//conditions:default": ["@capnp-cpp//src/kj:kj-async"], + }), +) + +# Build-graph hermeticity gate for the Rust I/O backend. `bazel build` this with +# --//:io_backend=rust to FAIL ANALYSIS if the workerd binary's transitive deps reach a +# forbidden concrete C++ I/O target (v1: the kj OS event loop / socket layer, kj-async-os), +# naming the offending dependency edge. In the default (cxx) config it only reports (those +# targets are the legitimate C++ backend there). This is the real compile/analysis-time +# replacement for the deleted grep gate; see //build:rust_io_backend.bzl. +rust_io_hermeticity( + name = "rust-io-hermeticity", + enforce = select({ + "//:io_backend_rust": True, + "//conditions:default": False, + }), + tags = ["manual"], + target = ":workerd", + visibility = ["//visibility:public"], +) + wd_cc_library( name = "alarm-scheduler", srcs = [ @@ -94,7 +137,7 @@ wd_cc_library( "//src/workerd/io", "//src/workerd/util:sqlite", "@capnp-cpp//src/kj", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", ], ) @@ -132,7 +175,7 @@ wd_cc_library( "//src/workerd/io:compatibility-date_capnp", "//src/workerd/jsg", "@capnp-cpp//src/kj", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", "@capnp-cpp//src/kj/compat:kj-gzip", "@capnp-cpp//src/kj/compat:kj-tls", ], @@ -149,6 +192,7 @@ wd_cc_library( visibility = ["//visibility:public"], deps = [ "//src/workerd/io", + "@capnp-cpp//src/capnp:capnp-rpc", ], ) @@ -243,7 +287,9 @@ wd_cc_library( "//src/workerd/io:worker-entrypoint", "//src/workerd/jsg", "//src/workerd/util:perfetto", + "//src/workerd/util:setup-async-io", "//src/workerd/util:websocket-error-handler", + "@capnp-cpp//src/capnp:capnp-rpc", "@capnp-cpp//src/kj/compat:kj-gzip", "@capnp-cpp//src/kj/compat:kj-tls", ], @@ -289,7 +335,7 @@ wd_cc_library( "@ada-url", "@capnp-cpp//src/capnp/compat:http-over-capnp", "@capnp-cpp//src/kj", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", "@capnp-cpp//src/kj/compat:kj-http", ], ) @@ -305,8 +351,9 @@ wd_cc_library( visibility = ["//visibility:public"], deps = [ ":workerd_capnp", + "//src/workerd/util:setup-async-io", "@capnp-cpp//src/kj", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", "@capnp-cpp//src/kj/compat:kj-http", ], ) @@ -415,9 +462,9 @@ kj_test( deps = [ ":fallback-service", ":workerd_capnp", + "//src/workerd/util:setup-async-io", "@capnp-cpp//src/capnp/compat:json", "@capnp-cpp//src/kj", - "@capnp-cpp//src/kj:kj-async", "@capnp-cpp//src/kj/compat:kj-http", ], ) diff --git a/src/workerd/server/cli-io-backend.c++ b/src/workerd/server/cli-io-backend.c++ new file mode 100644 index 00000000000..6416fc1b670 --- /dev/null +++ b/src/workerd/server/cli-io-backend.c++ @@ -0,0 +1,323 @@ +// Copyright (c) 2026 Cloudflare, Inc. +// Licensed under the Apache 2.0 license found in the LICENSE file or at: +// https://opensource.org/licenses/Apache-2.0 + +// Backend-specific implementations of workerd's CLI --watch file watcher and SIGTERM graceful +// drain (see cli-io-backend.h). The single WORKERD_RUST_IO_BACKEND_RUST #if that picks the native +// kj loop vs the tokio loop is confined to this TU, so workerd.c++'s call sites stay +// backend-agnostic -- exactly as //src/workerd/util:setup-async-io does for kj::setupAsyncIo(). + +#include "cli-io-backend.h" + +#if _WIN32 +#include +#include + +#if !WORKERD_RUST_IO_BACKEND_RUST +// --//:io_backend=rust ships kj-async-core only, which has no Win32EventPort and doesn't provide +// this header. The native watcher sites that need it are all gated out below. +#include +#endif +#include +#include +#else +#include +#include + +#include +#include + +#if !WORKERD_RUST_IO_BACKEND_RUST +// --//:io_backend=rust ships kj-async-core only, which has no UnixEventPort and doesn't provide +// this header. The native watcher + captureSignal/onSignal sites that need it are all gated out +// below; the tokio (kj-rs-io) loop covers those paths. +#include +#endif +#endif + +#if __linux__ +#include +#elif __APPLE__ || __FreeBSD__ || __OpenBSD__ || __NetBSD__ || __DragonFly__ +#define WORKERD_USE_KQUEUE_FOR_FILE_WATCHER 1 +#include +#include +#include +#endif + +#if WORKERD_RUST_IO_BACKEND_RUST +#include +#include +#endif + +#include +#include +#include + +namespace workerd::server { + +// The native (kj loop, kj::UnixEventPort::FdObserver) watcher is compiled only in the cxx config. +// Under --//:io_backend=rust there is no UnixEventPort, so --watch always uses the tokio watcher +// (below), which observes the same inotify/kqueue fds through tokio's AsyncFd. +#if !WORKERD_RUST_IO_BACKEND_RUST + +#if __linux__ + +// Class which uses inotify to watch a set of files and alert when they change. +class KjFileWatcher final: public FileWatcher { + public: + KjFileWatcher(kj::UnixEventPort& port) + : inotifyFd(makeInotify()), + observer(port, inotifyFd, kj::UnixEventPort::FdObserver::OBSERVE_READ) {} + + bool isSupported() override { + return true; + } + + void watch(kj::PathPtr path, kj::Maybe file) override { + // `file` is provided if available. The Linux implementation doesn't use it. + + auto pathStr = path.parent().toNativeString(true); + + int wd = watches.findOrCreate(pathStr, [&]() { + int wd; + uint32_t mask = IN_DELETE | IN_MODIFY | IN_MOVE | IN_CREATE; + KJ_SYSCALL(wd = inotify_add_watch(inotifyFd, pathStr.cStr(), mask)); + return decltype(watches)::Entry{kj::mv(pathStr), wd}; + }); + + auto& files = + filesWatched.findOrCreate(wd, [&]() { return decltype(filesWatched)::Entry{wd, {}}; }); + + files.upsert(kj::str(path.basename()[0]), [](auto&&...) {}); + } + + kj::Promise onChange() override { + kj::byte buffer[4096]{}; + + for (;;) { + ssize_t n; + KJ_NONBLOCKING_SYSCALL(n = read(inotifyFd, buffer, sizeof(buffer))); + + if (n < 0) { + // No more data to read. + co_await observer.whenBecomesReadable(); + continue; + } + + kj::byte* ptr = buffer; + while (n > 0) { + KJ_ASSERT(n >= sizeof(struct inotify_event)); + + auto& event = *reinterpret_cast(ptr); + size_t eventSize = sizeof(struct inotify_event) + event.len; + KJ_ASSERT(n >= eventSize); + KJ_ASSERT(eventSize % sizeof(void*) == 0); + ptr += eventSize; + n -= eventSize; + + if (event.len > 0 && event.name[0] != '\0') { + auto& watched = KJ_ASSERT_NONNULL(filesWatched.find(event.wd)); + if (watched.find(kj::StringPtr(event.name)) != kj::none) { + // HIT! We saw a change. + co_return; + } + } + } + } + } + + private: + kj::OwnFd inotifyFd; + kj::UnixEventPort::FdObserver observer; + + kj::HashMap watches; + kj::HashMap> filesWatched; + + static kj::OwnFd makeInotify() { + return KJ_SYSCALL_FD(inotify_init1(IN_NONBLOCK | IN_CLOEXEC)); + } +}; + +#elif WORKERD_USE_KQUEUE_FOR_FILE_WATCHER + +// Class which uses inotify to watch a set of files and alert when they change. +// +// This version uses kqueue to watch for changes in files. kqueue typically doesn't scale well +// to watching whole directory trees, since it must keep a file descriptor open for each watched +// file. However, for our use case, we don't really want to watch a directory tree anyway, we +// want to watch the specific set of files which were opened while parsing the config. This is +// not so bad, probably. +// +// Apple provides the FSEvents API as an alternative, but it seems way more complicated and I +// can't tell if it would provide a real advantage. Plus, kqueue works on BSD systems. +class KjFileWatcher final: public FileWatcher { + public: + KjFileWatcher(kj::UnixEventPort& port) + : kqueueFd(makeKqueue()), + observer(port, kqueueFd, kj::UnixEventPort::FdObserver::OBSERVE_READ) {} + + bool isSupported() override { + return true; + } + + void watch(kj::PathPtr path, kj::Maybe file) override { + KJ_IF_SOME(f, file) { + KJ_IF_SOME(fd, f.getFd()) { + // We need to duplicate the FD because the original will probably be closed later and + // closing the FD unregisters it from kqueue. + watchFd(KJ_SYSCALL_FD(dup(fd))); + return; + } + } + + // No existing file, open from disk. + watchFd(KJ_SYSCALL_FD(open(path.toNativeString(true).cStr(), O_RDONLY))); + } + + kj::Promise onChange() override { + for (;;) { + struct kevent event; + struct timespec timeout; + memset(&event, 0, sizeof(event)); + memset(&timeout, 0, sizeof(timeout)); + + int n; + KJ_SYSCALL(n = kevent(kqueueFd, nullptr, 0, &event, 1, &timeout)); + + if (n == 0) { + // No events, wait for the kqueue to become readable indicating an event has been + // delivered. + co_await observer.whenBecomesReadable(); + continue; + } else { + // We only pay attention to events that indicate changes in the first place, so there's + // no need to examine the event, it definitely means something changed. + co_return; + } + } + } + + private: + kj::OwnFd kqueueFd; + kj::UnixEventPort::FdObserver observer; + kj::Vector filesWatched; + + static kj::OwnFd makeKqueue() { + auto fd = KJ_SYSCALL_FD(kqueue()); + KJ_SYSCALL(fcntl(fd, F_SETFD, FD_CLOEXEC)); + return kj::mv(fd); + } + + void watchFd(kj::OwnFd fd) { + KJ_SYSCALL(fcntl(fd, F_SETFD, FD_CLOEXEC)); + + struct kevent change; + memset(&change, 0, sizeof(change)); + change.ident = fd.get(); + change.filter = EVFILT_VNODE; + change.flags = EV_ADD | EV_CLEAR; + change.fflags = NOTE_WRITE | NOTE_EXTEND | NOTE_DELETE | NOTE_RENAME; + KJ_SYSCALL(kevent(kqueueFd, &change, 1, nullptr, 0, nullptr)); + filesWatched.add(kj::mv(fd)); + } +}; + +#elif _WIN32 + +class KjFileWatcher final: public FileWatcher { + public: + KjFileWatcher(kj::Win32EventPort& port) {} + + bool isSupported() override { + return false; + } + + void watch(kj::PathPtr path, kj::Maybe file) override {} + + kj::Promise onChange() override { + return kj::NEVER_DONE; + } + + private: +}; + +#else + +// Dummy KjFileWatcher implementation for operating systems that aren't supported yet. +class KjFileWatcher final: public FileWatcher { + public: + KjFileWatcher(kj::UnixEventPort& port) {} + + bool isSupported() override { + return false; + } + + void watch(kj::PathPtr path, kj::Maybe file) override {} + + kj::Promise onChange() override { + return kj::NEVER_DONE; + } + + private: +}; + +#endif // #__linux__, #else + +#endif // !WORKERD_RUST_IO_BACKEND_RUST + +// FileWatcher for --//:io_backend=rust: the same inotify/kqueue backends as the native watcher +// (line-for-line ports living in src/rust/cxx/kj-rs-io), but fd readiness is awaited through +// tokio's AsyncFd instead of kj::UnixEventPort::FdObserver, which doesn't exist on the tokio +// loop. All fds it creates are CLOEXEC, so reloadFromConfigChange()'s execve() doesn't leak +// them (the deliberately-inherited sockets get FIONCLEX'd there explicitly). +#if WORKERD_RUST_IO_BACKEND_RUST +class TokioFileWatcher final: public FileWatcher { + public: + bool isSupported() override { + return inner.isSupported(); + } + + void watch(kj::PathPtr path, kj::Maybe file) override { + inner.watch(path, file); + } + + kj::Promise onChange() override { + return inner.onChange(); + } + + private: + kj_rs_io::FileWatcher inner; +}; +#endif // WORKERD_RUST_IO_BACKEND_RUST + +kj::Own makeFileWatcher([[maybe_unused]] kj::AsyncIoContext& io) { +#if WORKERD_RUST_IO_BACKEND_RUST + // No UnixEventPort exists on the tokio loop; use the AsyncFd-backed watcher instead. + return kj::heap(); +#elif _WIN32 + return kj::heap(io.win32EventPort); +#else + return kj::heap(io.unixEventPort); +#endif +} + +#if !_WIN32 +void captureSigterm() { +#if !WORKERD_RUST_IO_BACKEND_RUST + kj::UnixEventPort::captureSignal(SIGTERM); +#endif + // Under --//:io_backend=rust this is a no-op: tokio's signal driver watches SIGTERM instead, and + // capturing it here would block the signal in the thread's mask and prevent the tokio handler + // from ever being invoked (there is no UnixEventPort under that backend anyway). +} + +kj::Promise onSigterm([[maybe_unused]] kj::AsyncIoContext& io) { +#if WORKERD_RUST_IO_BACKEND_RUST + return kj_rs_io::onSignal(SIGTERM); +#else + return io.unixEventPort.onSignal(SIGTERM).ignoreResult(); +#endif +} +#endif // !_WIN32 + +} // namespace workerd::server diff --git a/src/workerd/server/cli-io-backend.h b/src/workerd/server/cli-io-backend.h new file mode 100644 index 00000000000..bbabf41993c --- /dev/null +++ b/src/workerd/server/cli-io-backend.h @@ -0,0 +1,68 @@ +// Copyright (c) 2026 Cloudflare, Inc. +// Licensed under the Apache 2.0 license found in the LICENSE file or at: +// https://opensource.org/licenses/Apache-2.0 + +#pragma once + +// The two pieces of workerd's CLI main() whose implementation differs between the native kj event +// loop (default, cxx build) and the tokio loop (--//:io_backend=rust): the --watch file watcher +// and the SIGTERM graceful-drain signal. Both are presented here as backend-agnostic entry points +// so that workerd.c++ links against them unchanged in both configs -- no +// WORKERD_RUST_IO_BACKEND_RUST at the call site. The single compile-time #if that selects the +// backend lives in cli-io-backend.c++, mirroring how //src/workerd/util:setup-async-io makes +// kj::setupAsyncIo() transparent. + +#include +#include +#include +#include + +namespace workerd::server { + +// Interface for watching the files the server depends on (parsed config files, worker source +// files, and the server binary itself) and alerting when any of them change; drives --watch. Two +// implementations, selected by the active I/O backend inside makeFileWatcher(): +// +// - the native watcher: inotify (Linux) / kqueue (macOS, BSDs) readiness observed through +// kj::UnixEventPort::FdObserver. Requires the native KJ event loop. +// - the tokio watcher: same platform backends, but readiness observed through tokio's AsyncFd +// (kj_rs_io::FileWatcher from src/rust/cxx/kj-rs-io); used under --//:io_backend=rust, where no +// UnixEventPort exists. +// +// Everything downstream (SchemaFileImpl's watch registration, waitForChanges()'s coalescing, +// serveImpl()'s re-exec loop) is shared between the two. +class FileWatcher { + public: + virtual ~FileWatcher() noexcept(false) = default; + + // False on platforms where watching is not implemented (callers report a CLI error). + virtual bool isSupported() = 0; + + // Adds `path` to the watched set. `file` is an already-open handle for the same path, if + // available (the kqueue backends watch the open file directly; others open by path). + virtual void watch(kj::PathPtr path, kj::Maybe file) = 0; + + // Resolves the next time any watched file changes. Changes are queued by the kernel, not + // lost between calls; call again after resolution to wait for further changes. + virtual kj::Promise onChange() = 0; +}; + +// Constructs the FileWatcher appropriate for the active I/O backend. Under the native (cxx) loop +// this is the UnixEventPort::FdObserver-based watcher, driven by `io`'s event port; under +// --//:io_backend=rust it is the tokio AsyncFd-based watcher, which ignores `io`. +kj::Own makeFileWatcher(kj::AsyncIoContext& io); + +#if !_WIN32 +// Captures SIGTERM so the native loop can later deliver it to onSigterm(). Under the native (cxx) +// loop this is kj::UnixEventPort::captureSignal(SIGTERM); under --//:io_backend=rust it is a no-op +// (tokio's own signal driver handles SIGTERM, and capturing would block it from that driver). Call +// once in main() before the event loop is created. +void captureSigterm(); + +// Resolves when SIGTERM is received; used as Server::run()'s drainWhen promise. Under the native +// (cxx) loop this is io.unixEventPort.onSignal(SIGTERM); under --//:io_backend=rust it is a +// tokio-signal-backed promise from kj-rs-io (and `io` is ignored). +kj::Promise onSigterm(kj::AsyncIoContext& io); +#endif + +} // namespace workerd::server diff --git a/src/workerd/server/workerd.c++ b/src/workerd/server/workerd.c++ index 337a7627c80..cf7a9d265b6 100644 --- a/src/workerd/server/workerd.c++ +++ b/src/workerd/server/workerd.c++ @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -48,17 +49,6 @@ #include #include #include - -#include -#endif - -#if __linux__ -#include -#elif __APPLE__ || __FreeBSD__ || __OpenBSD__ || __NetBSD__ || __DragonFly__ -#define WORKERD_USE_KQUEUE_FOR_FILE_WATCHER 1 -#include -#include -#include #endif #ifdef __GLIBC__ @@ -170,211 +160,6 @@ constexpr capnp::ReaderOptions CONFIG_READER_OPTIONS = { // ======================================================================================= -#if __linux__ - -// Class which uses inotify to watch a set of files and alert when they change. -class FileWatcher { - public: - FileWatcher(kj::UnixEventPort& port) - : inotifyFd(makeInotify()), - observer(port, inotifyFd, kj::UnixEventPort::FdObserver::OBSERVE_READ) {} - - bool isSupported() { - return true; - } - - void watch(kj::PathPtr path, kj::Maybe file) { - // `file` is provided if available. The Linux implementation doesn't use it. - - auto pathStr = path.parent().toNativeString(true); - - int wd = watches.findOrCreate(pathStr, [&]() { - int wd; - uint32_t mask = IN_DELETE | IN_MODIFY | IN_MOVE | IN_CREATE; - KJ_SYSCALL(wd = inotify_add_watch(inotifyFd, pathStr.cStr(), mask)); - return decltype(watches)::Entry{kj::mv(pathStr), wd}; - }); - - auto& files = - filesWatched.findOrCreate(wd, [&]() { return decltype(filesWatched)::Entry{wd, {}}; }); - - files.upsert(kj::str(path.basename()[0]), [](auto&&...) {}); - } - - kj::Promise onChange() { - kj::byte buffer[4096]{}; - - for (;;) { - ssize_t n; - KJ_NONBLOCKING_SYSCALL(n = read(inotifyFd, buffer, sizeof(buffer))); - - if (n < 0) { - // No more data to read. - co_await observer.whenBecomesReadable(); - continue; - } - - kj::byte* ptr = buffer; - while (n > 0) { - KJ_ASSERT(n >= sizeof(struct inotify_event)); - - auto& event = *reinterpret_cast(ptr); - size_t eventSize = sizeof(struct inotify_event) + event.len; - KJ_ASSERT(n >= eventSize); - KJ_ASSERT(eventSize % sizeof(void*) == 0); - ptr += eventSize; - n -= eventSize; - - if (event.len > 0 && event.name[0] != '\0') { - auto& watched = KJ_ASSERT_NONNULL(filesWatched.find(event.wd)); - if (watched.find(kj::StringPtr(event.name)) != kj::none) { - // HIT! We saw a change. - co_return; - } - } - } - } - } - - private: - kj::OwnFd inotifyFd; - kj::UnixEventPort::FdObserver observer; - - kj::HashMap watches; - kj::HashMap> filesWatched; - - static kj::OwnFd makeInotify() { - return KJ_SYSCALL_FD(inotify_init1(IN_NONBLOCK | IN_CLOEXEC)); - } -}; - -#elif WORKERD_USE_KQUEUE_FOR_FILE_WATCHER - -// Class which uses inotify to watch a set of files and alert when they change. -// -// This version uses kqueue to watch for changes in files. kqueue typically doesn't scale well -// to watching whole directory trees, since it must keep a file descriptor open for each watched -// file. However, for our use case, we don't really want to watch a directory tree anyway, we -// want to watch the specific set of files which were opened while parsing the config. This is -// not so bad, probably. -// -// Apple provides the FSEvents API as an alternative, but it seems way more complicated and I -// can't tell if it would provide a real advantage. Plus, kqueue works on BSD systems. -class FileWatcher { - public: - FileWatcher(kj::UnixEventPort& port) - : kqueueFd(makeKqueue()), - observer(port, kqueueFd, kj::UnixEventPort::FdObserver::OBSERVE_READ) {} - - bool isSupported() { - return true; - } - - void watch(kj::PathPtr path, kj::Maybe file) { - KJ_IF_SOME(f, file) { - KJ_IF_SOME(fd, f.getFd()) { - // We need to duplicate the FD because the original will probably be closed later and - // closing the FD unregisters it from kqueue. - watchFd(KJ_SYSCALL_FD(dup(fd))); - return; - } - } - - // No existing file, open from disk. - watchFd(KJ_SYSCALL_FD(open(path.toNativeString(true).cStr(), O_RDONLY))); - } - - kj::Promise onChange() { - for (;;) { - struct kevent event; - struct timespec timeout; - memset(&event, 0, sizeof(event)); - memset(&timeout, 0, sizeof(timeout)); - - int n; - KJ_SYSCALL(n = kevent(kqueueFd, nullptr, 0, &event, 1, &timeout)); - - if (n == 0) { - // No events, wait for the kqueue to become readable indicating an event has been - // delivered. - co_await observer.whenBecomesReadable(); - continue; - } else { - // We only pay attention to events that indicate changes in the first place, so there's - // no need to examine the event, it definitely means something changed. - co_return; - } - } - } - - private: - kj::OwnFd kqueueFd; - kj::UnixEventPort::FdObserver observer; - kj::Vector filesWatched; - - static kj::OwnFd makeKqueue() { - auto fd = KJ_SYSCALL_FD(kqueue()); - KJ_SYSCALL(fcntl(fd, F_SETFD, FD_CLOEXEC)); - return kj::mv(fd); - } - - void watchFd(kj::OwnFd fd) { - KJ_SYSCALL(fcntl(fd, F_SETFD, FD_CLOEXEC)); - - struct kevent change; - memset(&change, 0, sizeof(change)); - change.ident = fd.get(); - change.filter = EVFILT_VNODE; - change.flags = EV_ADD | EV_CLEAR; - change.fflags = NOTE_WRITE | NOTE_EXTEND | NOTE_DELETE | NOTE_RENAME; - KJ_SYSCALL(kevent(kqueueFd, &change, 1, nullptr, 0, nullptr)); - filesWatched.add(kj::mv(fd)); - } -}; - -#elif _WIN32 - -class FileWatcher { - public: - FileWatcher(kj::Win32EventPort& port) {} - - bool isSupported() { - return false; - } - - void watch(kj::PathPtr path, kj::Maybe file) {} - - kj::Promise onChange() { - return kj::NEVER_DONE; - } - - private: -}; - -#else - -// Dummy FileWatcher implementation for operating systems that aren't supported yet. -class FileWatcher { - public: - FileWatcher(kj::UnixEventPort& port) {} - - bool isSupported() { - return false; - } - - void watch(kj::PathPtr path, kj::Maybe file) {} - - kj::Promise onChange() { - return kj::NEVER_DONE; - } - - private: -}; - -#endif // #__linux__, #else - -// ======================================================================================= - kj::Maybe> tryImportBulitin(kj::StringPtr name); // Callbacks for capnp::SchemaFileLoader. Implementing this interface lets us control import @@ -1158,11 +943,7 @@ class CliMain final: public SchemaFileImpl::ErrorReporter { } void watch() { -#if _WIN32 - auto& w = watcher.emplace(io.win32EventPort); -#else - auto& w = watcher.emplace(io.unixEventPort); -#endif + FileWatcher& w = *watcher.emplace(makeFileWatcher(io)); if (!w.isSupported()) { CLI_ERROR("File watching is not yet implemented on your OS. Sorry! Pull requests welcome!"); } @@ -1217,7 +998,8 @@ class CliMain final: public SchemaFileImpl::ErrorReporter { schemaParser.loadCompiledTypeAndDependencies(); parsedSchema = schemaParser.parseFile(kj::heap(fs->getRoot(), - fs->getCurrentPath(), kj::mv(path), nullptr, importPath, kj::mv(file), watcher, *this)); + fs->getCurrentPath(), kj::mv(path), nullptr, importPath, kj::mv(file), + watcher.map([](kj::Own& w) -> FileWatcher& { return *w; }), *this)); // Construct a list of top-level constants of type `Config`. If there is exactly one, // we can use it by default. @@ -1410,7 +1192,7 @@ class CliMain final: public SchemaFileImpl::ErrorReporter { // someone to fix the config. context.warning( "Can't start server due to config errors, waiting for config files to change..."); - waitForChanges(w).wait(io.waitScope); + waitForChanges(*w).wait(io.waitScope); reloadFromConfigChange(); } else { // Errors were reported earlier, so context.exit() will exit with a non-zero status. @@ -1439,7 +1221,7 @@ class CliMain final: public SchemaFileImpl::ErrorReporter { KJ_MAP(flag, config.getV8Flags()) -> kj::StringPtr { return flag; }, platform.get()); auto promise = func(v8System, config); KJ_IF_SOME(w, watcher) { - promise = promise.exclusiveJoin(waitForChanges(w).then([this]() { + promise = promise.exclusiveJoin(waitForChanges(*w).then([this]() { // Watch succeeded. reloadFromConfigChange(); })); @@ -1467,8 +1249,8 @@ class CliMain final: public SchemaFileImpl::ErrorReporter { return server->run(v8System, config); #else return server->run(v8System, config, - // Gracefully drain when SIGTERM is received. - io.unixEventPort.onSignal(SIGTERM).ignoreResult()); + // Gracefully drain when SIGTERM is received (backend-specific; see cli-io-backend.h). + onSigterm(io)); #endif }); } @@ -1563,7 +1345,7 @@ class CliMain final: public SchemaFileImpl::ErrorReporter { bool gcStress = false; bool allAutogates = false; kj::Maybe testCompatDate; - kj::Maybe watcher; + kj::Maybe> watcher; kj::Own fs = kj::newDiskFilesystem(); kj::AsyncIoContext io = kj::setupAsyncIo(); @@ -1757,7 +1539,7 @@ int main(int argc, char* argv[]) { workerd::server::StructuredLoggingProcessContext context(argv[0]); #if !_WIN32 - kj::UnixEventPort::captureSignal(SIGTERM); + workerd::server::captureSigterm(); #endif workerd::server::CliMain mainObject(context, argv); diff --git a/src/workerd/tests/BUILD.bazel b/src/workerd/tests/BUILD.bazel index 298440ae427..3658b0d4c1a 100644 --- a/src/workerd/tests/BUILD.bazel +++ b/src/workerd/tests/BUILD.bazel @@ -39,6 +39,7 @@ wd_cc_library( "//src/workerd/jsg", "//src/workerd/server:workerd-api", "//src/workerd/util:autogate", + "//src/workerd/util:setup-async-io", ], ) diff --git a/src/workerd/util/BUILD.bazel b/src/workerd/util/BUILD.bazel index b902f621d26..62ccc6b5229 100644 --- a/src/workerd/util/BUILD.bazel +++ b/src/workerd/util/BUILD.bazel @@ -1,6 +1,7 @@ load("@bazel_skylib//lib:selects.bzl", "selects") load("@bazel_skylib//rules:common_settings.bzl", "bool_flag") load("//:build/kj_test.bzl", "kj_test") +load("//:build/rust_io_backend.bzl", "rust_io_backend_local_defines") load("//:build/wd_cc_library.bzl", "wd_cc_library") bool_flag( @@ -87,7 +88,7 @@ wd_cc_library( deps = [ ":duration-exceeded-logger", "@capnp-cpp//src/kj", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", # TODO(cleanup): Only for abortable.h, factor out "@capnp-cpp//src/kj/compat:kj-http", ], @@ -135,7 +136,7 @@ wd_cc_library( ":strings", "//src/workerd/jsg:memory-tracker", "@capnp-cpp//src/kj", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", ], ) @@ -220,7 +221,7 @@ wd_cc_library( ":account-limits", ":sentry", ":sqlite-metering", - "@capnp-cpp//src/kj:kj-async", + "@capnp-cpp//src/kj:kj-async-core", ], ) @@ -465,3 +466,33 @@ kj_test( src = "state-machine-test.c++", deps = [":state-machine"], ) + +# The kj::setupAsyncIo() "seam": supplies that single symbol appropriately for each I/O backend, +# so call sites depend on this target and call kj::setupAsyncIo() unchanged in both configs (no +# #if at the call site, and kj/async-io.h itself is untouched). +# * default (cxx): setup-async-io-tokio.c++ compiles to an empty TU; native kj::setupAsyncIo() +# comes from the :kj-async umbrella. Byte-identical to depending on :kj-async directly. +# * --//:io_backend=rust: kj-async-os (the OS event loop) is NOT linked; the .c++ (gated on +# WORKERD_RUST_IO_BACKEND_RUST) defines a tokio-backed kj::setupAsyncIo() plus an inert +# kj::UnixEventPort (whose real defs live in the unlinked async-unix.c++). //deps:rust_runtime +# supplies the Rust allocator shims. +wd_cc_library( + name = "setup-async-io", + srcs = ["setup-async-io-tokio.c++"], + local_defines = rust_io_backend_local_defines(), + visibility = ["//visibility:public"], + deps = select({ + "//:io_backend_rust": [ + "//deps:rust_runtime", + # The abstract async I/O layer (kj/async-io.h: AsyncIoStream / Network / AsyncIoContext) + # the shim's setupAsyncIo() is written against... + "@capnp-cpp//src/kj:kj-async-io", + # ...plus the declaration-only view of the OS event port (kj/async-unix.h), whose + # out-of-line member definitions the shim supplies instead of :kj-async-os. + "@capnp-cpp//src/kj:kj-async-os-hdrs", + "//src/rust/cxx/kj-rs-io", + "//src/rust/cxx/kj-rs-io:kj-rs-io-lib", + ], + "//conditions:default": ["@capnp-cpp//src/kj:kj-async"], + }), +) diff --git a/src/workerd/util/setup-async-io-tokio.c++ b/src/workerd/util/setup-async-io-tokio.c++ new file mode 100644 index 00000000000..23aab680196 --- /dev/null +++ b/src/workerd/util/setup-async-io-tokio.c++ @@ -0,0 +1,173 @@ +// Copyright (c) 2026 Cloudflare, Inc. +// Licensed under the Apache 2.0 license found in the LICENSE file or at: +// https://opensource.org/licenses/Apache-2.0 + +// Provides the kj::setupAsyncIo() *symbol* for the rust I/O backend, so every existing +// kj::setupAsyncIo() call site links unchanged under --//:io_backend=rust -- no per-call-site +// #if, no divergence from upstream kj (kj/async-io.h is untouched). +// +// Under --//:io_backend=rust the native OS event loop (kj-async-os, which defines both +// kj::setupAsyncIo() and kj::UnixEventPort's member functions) is NOT linked. This TU supplies a +// tokio-backed kj::setupAsyncIo() that repackages src/rust/cxx/kj-rs-io's setupTokioAsyncIo() into a +// kj::AsyncIoContext. +// +// kj::AsyncIoContext still declares `UnixEventPort& unixEventPort` (a concrete reference we must +// not change -- it's upstream kj public API). kj::UnixEventPort is `final`, but that only forbids +// subclassing; its method *definitions* live in async-unix.c++, which is absent from this link. +// So we define an inert kj::UnixEventPort here (its ctor/dtor + the EventPort/SleepHooks virtuals) +// with no ODR competitor, construct one, and bind the reference to it. It is never driven -- the +// event loop runs on kj_rs_tokio::TokioEventPort; signals use kj_rs_io::onSignal() -- and nothing +// reads AsyncIoContext::unixEventPort under the rust backend (verified: the only readers, SIGTERM +// drain + --watch, are #if !WORKERD_RUST_IO_BACKEND_RUST). Its methods KJ_UNIMPLEMENTED as a +// backstop: if anything ever does drive it, the link/run fails loudly rather than silently. +// +// The whole TU is gated on WORKERD_RUST_IO_BACKEND_RUST so that in the default (cxx) build it is +// empty and cannot ODR-clash with kj-async-os's real definitions. +// +// In the rust config the reverse hazard exists: if kj-async-os is ever accidentally linked back +// in, this TU's definitions collide with the native ones, and with static archives the winner is +// LINK-ORDER-DEPENDENT -- a duplicate-symbol error if you are lucky, the native setupAsyncIo +// silently winning (wrong event loop) or pairing with this inert UnixEventPort +// (KJ_UNIMPLEMENTED at runtime) if you are not. The guard against that is the build-graph gate +// //src/workerd/server:rust-io-hermeticity (build/rust_io_backend.bzl), which is tagged `manual`: +// it is NOT implied by building :workerd and must be built explicitly in the rust-config CI lane. + +#if WORKERD_RUST_IO_BACKEND_RUST + +#include + +#include +#if _WIN32 +#include +#else +#include +#endif +#include + +namespace kj { + +#if _WIN32 + +// ---- Inert Win32EventPort ---------------------------------------------------------------------- +// On Windows kj::AsyncIoContext declares `Win32EventPort& win32EventPort`. Unlike UnixEventPort, +// Win32EventPort is an abstract interface, so instead of supplying out-of-line definitions for an +// unlinked concrete class we define our own inert implementation. Same contract as the unix inert +// port below: never driven, KJ_UNIMPLEMENTED as a backstop. +class InertWin32EventPort final: public Win32EventPort { + public: + InertWin32EventPort(): clock(systemPreciseMonotonicClock()), timerImpl(clock.now()) {} + + bool wait() override { + KJ_UNIMPLEMENTED("Win32EventPort is inert under --//:io_backend=rust (tokio drives the loop)"); + } + bool poll() override { + KJ_UNIMPLEMENTED("Win32EventPort is inert under --//:io_backend=rust (tokio drives the loop)"); + } + void wake() const override { + KJ_UNIMPLEMENTED("Win32EventPort is inert under --//:io_backend=rust (tokio drives the loop)"); + } + Own observeIo(HANDLE handle) override { + KJ_UNIMPLEMENTED("Win32EventPort is inert under --//:io_backend=rust (tokio drives the loop)"); + } + Own observeSignalState(HANDLE handle) override { + KJ_UNIMPLEMENTED("Win32EventPort is inert under --//:io_backend=rust (tokio drives the loop)"); + } + void allowApc() override { + KJ_UNIMPLEMENTED("Win32EventPort is inert under --//:io_backend=rust (tokio drives the loop)"); + } + Timer& getTimer() override { + return timerImpl; + } + + private: + const MonotonicClock& clock; + TimerImpl timerImpl; +}; + +using InertEventPort = InertWin32EventPort; + +#else // _WIN32 + +// ---- Inert kj::UnixEventPort (real definitions live in the unlinked async-unix.c++) ----------- + +#if !KJ_USE_KQUEUE +// On epoll platforms UnixEventPort holds a `Maybe>` whose type is only +// forward-declared in async-unix.h; the real definition lives in the unlinked async-unix.c++. +// Defining ~UnixEventPort() below requires the type to be complete. The inert port never +// constructs one (the Maybe stays none), so an empty definition satisfies the compiler and the +// dispose path is never reached. +struct UnixEventPort::ChildSet {}; +#endif + +UnixEventPort::UnixEventPort(): clock(systemPreciseMonotonicClock()), timerImpl(clock.now()) {} + +UnixEventPort::~UnixEventPort() noexcept(false) {} + +bool UnixEventPort::wait() { + KJ_UNIMPLEMENTED("UnixEventPort is inert under --//:io_backend=rust (tokio drives the loop)"); +} + +bool UnixEventPort::poll() { + KJ_UNIMPLEMENTED("UnixEventPort is inert under --//:io_backend=rust (tokio drives the loop)"); +} + +void UnixEventPort::wake() const { + KJ_UNIMPLEMENTED("UnixEventPort is inert under --//:io_backend=rust (tokio drives the loop)"); +} + +#if KJ_USE_EPOLL +// This guard mirrors kj/async-unix.h, which declares UnixEventPort::setRunnable only under +// KJ_USE_EPOLL. On kqueue platforms the member does not exist (defining it would not compile); +// there UnixEventPort inherits kj::EventPort::setRunnable, whose default is a no-op — harmless, +// since this inert port is never installed as an EventLoop's port. +void UnixEventPort::setRunnable(bool runnable) { + KJ_UNIMPLEMENTED("UnixEventPort is inert under --//:io_backend=rust (tokio drives the loop)"); +} +#endif + +void UnixEventPort::updateNextTimerEvent(kj::Maybe time) {} + +kj::TimePoint UnixEventPort::getTimeWhileSleeping() { + KJ_UNIMPLEMENTED("UnixEventPort is inert under --//:io_backend=rust (tokio drives the loop)"); +} + +using InertEventPort = UnixEventPort; + +#endif // _WIN32, !_WIN32 + +// ---- The tokio-backed kj::setupAsyncIo() ------------------------------------------------------ + +AsyncIoContext setupAsyncIo(kj::Maybe observer) { + // The observer hook is a native-EventLoop concept the tokio loop does not surface; callers + // under --//:io_backend=rust pass kj::none. + struct Holder { + kj_rs_io::TokioAsyncIoContext tokio; + InertEventPort inertPort; + Holder(): tokio(kj_rs_io::setupTokioAsyncIo()) {} + }; + auto holder = kj::heap(); + + auto& lowLevel = holder->tokio.getLowLevelProvider(); + auto& provider = holder->tokio.getProvider(); + auto& waitScope = holder->tokio.getWaitScope(); + auto& inertPort = holder->inertPort; + + // Non-owning handles into the heap Holder; the Holder is attached to the lowLevelProvider handle + // so it (and thus the tokio context + inert port) is torn down exactly once when the returned + // AsyncIoContext's lowLevelProvider is destroyed. The provider handle uses NullDisposer so + // AsyncIoContext's earlier destruction of `provider` is a no-op. + // + // Lifetime of the returned references: the Holder lives inside AsyncIoContext::lowLevelProvider, + // so `waitScope` and `unixEventPort` (bound to holder->inertPort) remain valid for the full + // lifetime of the AsyncIoContext. They dangle only once the AsyncIoContext itself is destroyed, + // same as with the native kj::setupAsyncIo(). + kj::Own lowLevelOwn(&lowLevel, kj::NullDisposer::instance); + kj::Own providerOwn(&provider, kj::NullDisposer::instance); + lowLevelOwn = lowLevelOwn.attach(kj::mv(holder)); + + return AsyncIoContext{kj::mv(lowLevelOwn), kj::mv(providerOwn), waitScope, inertPort}; +} + +} // namespace kj + +#endif // WORKERD_RUST_IO_BACKEND_RUST