This is an automated email from the ASF dual-hosted git repository.
hubcio pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iggy.git
The following commit(s) were added to refs/heads/master by this push:
new f77d64b2f fix(server): distribute TCP-TLS and WSS connections across
shards (#4193)
f77d64b2f is described below
commit f77d64b2faf01c11e77242c1a073b6b9b200d7aa
Author: Hubert Gruszecki <[email protected]>
AuthorDate: Tue Sep 15 16:07:08 2026 +0200
fix(server): distribute TCP-TLS and WSS connections across shards (#4193)
---
core/integration/tests/server/specific.rs | 187 ++++++++++++-
core/message_bus/src/client_listener/mod.rs | 31 ++-
core/message_bus/src/client_listener/tcp_tls.rs | 27 +-
core/message_bus/src/client_listener/wss.rs | 16 +-
core/message_bus/src/fd_transfer.rs | 26 +-
core/message_bus/src/installer/common.rs | 7 +-
core/message_bus/src/installer/mod.rs | 56 +++-
core/message_bus/src/installer/replica.rs | 14 +-
core/message_bus/src/installer/tcp.rs | 5 +
core/message_bus/src/installer/tcp_tls.rs | 9 +-
core/message_bus/src/installer/wss.rs | 5 +-
core/message_bus/src/lib.rs | 52 ++--
.../src/lifecycle/connection_registry.rs | 11 +-
core/message_bus/src/replica/io.rs | 45 ++-
core/message_bus/src/transports/tcp_tls.rs | 6 +-
core/message_bus/src/transports/wss.rs | 28 +-
core/message_bus/tests/common/mod.rs | 111 +++++++-
core/message_bus/tests/tcp_tls_client_listener.rs | 5 +
core/message_bus/tests/wss_client_listener.rs | 108 ++++++++
core/server/src/boot/listeners.rs | 64 ++---
core/server/src/boot/mod.rs | 7 +-
core/server/src/dispatch/mod.rs | 11 +-
core/server/src/dispatch/test_support.rs | 18 +-
core/shard/src/coordinator.rs | 308 +++++++++++++++++++--
core/shard/src/lib.rs | 23 +-
core/shard/src/router.rs | 24 ++
core/simulator/src/bus.rs | 22 +-
27 files changed, 956 insertions(+), 270 deletions(-)
diff --git a/core/integration/tests/server/specific.rs
b/core/integration/tests/server/specific.rs
index 5bf367083..3e3281d72 100644
--- a/core/integration/tests/server/specific.rs
+++ b/core/integration/tests/server/specific.rs
@@ -20,33 +20,70 @@ use
crate::server::scenarios::{reconnect_after_restart_scenario, restart_offset_
use crate::server::scenarios::{
segment_rotation_race_scenario, tcp_tls_scenario, websocket_tls_scenario,
};
+use iggy::prelude::*;
+use integration::harness::TestHarness;
use integration::iggy_harness;
+use std::collections::HashSet;
+use std::time::Duration;
+use tokio::io::AsyncWriteExt;
+use tokio::net::TcpStream;
+
+const TLS_CLIENTS_PER_SHARD: usize = 2;
+const TLS_TEST_TIMEOUT: Duration = Duration::from_secs(10);
+const TLS_POLL_INTERVAL: Duration = Duration::from_millis(20);
#[iggy_harness(
test_client_transport = TcpTlsGenerated,
- server(tls = generated)
+ server(tls = generated, logging.level = "info")
)]
async fn tcp_tls_scenario_should_be_valid(harness: &TestHarness) {
- let client = harness.root_client().await.unwrap();
- tcp_tls_scenario::run(&client).await;
+ let shard_count = tls_test_shard_count(harness).await;
+ let slow_clients = stalled_tls_clients(harness, shard_count).await;
+ let clients = harness
+ .root_clients(shard_count * TLS_CLIENTS_PER_SHARD)
+ .await
+ .unwrap();
+ for client in &clients {
+ tcp_tls_scenario::run(client).await;
+ }
+ drop(slow_clients);
+ assert_tls_client_sharding_and_cleanup(harness, &clients, "TCP-TLS",
shard_count).await;
}
#[iggy_harness(
test_client_transport = TcpTlsSelfSigned,
- server(tls = self_signed)
+ server(tls = self_signed, logging.level = "info")
)]
async fn tcp_tls_self_signed_scenario_should_be_valid(harness: &TestHarness) {
- let client = harness.root_client().await.unwrap();
- tcp_tls_scenario::run(&client).await;
+ let shard_count = tls_test_shard_count(harness).await;
+ let slow_clients = stalled_tls_clients(harness, shard_count).await;
+ let clients = harness
+ .root_clients(shard_count * TLS_CLIENTS_PER_SHARD)
+ .await
+ .unwrap();
+ for client in &clients {
+ tcp_tls_scenario::run(client).await;
+ }
+ drop(slow_clients);
+ assert_tls_client_sharding_and_cleanup(harness, &clients, "TCP-TLS",
shard_count).await;
}
#[iggy_harness(
test_client_transport = WebSocketTlsGenerated,
- server(websocket_tls = generated)
+ server(websocket_tls = generated, logging.level = "info")
)]
async fn websocket_tls_scenario_should_be_valid(harness: &TestHarness) {
- let client = harness.root_client().await.unwrap();
- websocket_tls_scenario::run(&client).await;
+ let shard_count = tls_test_shard_count(harness).await;
+ let slow_clients = stalled_tls_clients(harness, shard_count).await;
+ let clients = harness
+ .root_clients(shard_count * TLS_CLIENTS_PER_SHARD)
+ .await
+ .unwrap();
+ for client in &clients {
+ websocket_tls_scenario::run(client).await;
+ }
+ drop(slow_clients);
+ assert_tls_client_sharding_and_cleanup(harness, &clients, "WSS",
shard_count).await;
}
#[iggy_harness]
@@ -141,3 +178,135 @@ async fn restart_offset_skip(harness: &mut TestHarness) {
async fn segment_rotation_scenario(harness: &TestHarness) {
segment_rotation_race_scenario::run(harness).await;
}
+
+async fn tls_test_shard_count(harness: &TestHarness) -> usize {
+ tokio::time::timeout(TLS_TEST_TIMEOUT, async {
+ loop {
+ let logs = harness.server().stdout_plain();
+ if let Some(line) = logs
+ .lines()
+ .find(|line| line.contains("server bootstrap dispatched;
awaiting shard runtimes"))
+ {
+ let shard_count = line
+ .split_whitespace()
+ .find_map(|field| field.strip_prefix("shards_count="))
+ .expect("bootstrap log must carry shards_count")
+ .parse::<usize>()
+ .expect("bootstrap shards_count must be numeric");
+ assert!(shard_count > 0, "server must start at least one
shard");
+ return shard_count;
+ }
+ tokio::time::sleep(TLS_POLL_INTERVAL).await;
+ }
+ })
+ .await
+ .expect("server must report its resolved shard count")
+}
+
+async fn assert_tls_client_sharding_and_cleanup(
+ harness: &TestHarness,
+ clients: &[IggyClient],
+ transport: &str,
+ shard_count: usize,
+) {
+ assert_tls_client_sharding(harness, clients, transport, shard_count).await;
+ for client in clients {
+ client.disconnect().await.unwrap();
+ }
+
+ // Each fresh observer has no partition connection and lives on a different
+ // shard. Seeing every observer proves that the gather includes every
shard.
+ let observers = harness.root_clients(shard_count).await.unwrap();
+ let observer_ids =
+ assert_tls_client_sharding(harness, &observers, transport,
shard_count).await;
+ tokio::time::timeout(TLS_TEST_TIMEOUT, async {
+ loop {
+ let connected = observers[0].get_clients().await.unwrap();
+ let connected_ids: HashSet<_> =
+ connected.iter().map(|client| client.client_id).collect();
+ if connected.len() == observers.len() && connected_ids ==
observer_ids {
+ break;
+ }
+ tokio::time::sleep(TLS_POLL_INTERVAL).await;
+ }
+ })
+ .await
+ .expect("every shard must report only its observer after client
disconnect");
+ for observer in observers {
+ observer.disconnect().await.unwrap();
+ }
+}
+
+async fn assert_tls_client_sharding(
+ harness: &TestHarness,
+ clients: &[IggyClient],
+ transport: &str,
+ shard_count: usize,
+) -> HashSet<u32> {
+ let mut client_ids = HashSet::with_capacity(clients.len());
+ for client in clients {
+ client_ids.insert(client.get_me().await.unwrap().client_id);
+ }
+ assert_eq!(client_ids.len(), clients.len(), "client IDs must be unique");
+
+ // The wire exposes only the sequence tail. Match it to the router's
+ // full ID and thread name to prove execution placement after handoff.
+ let install_marker = format!("installing delegated {transport} client fd");
+ tokio::time::timeout(TLS_TEST_TIMEOUT, async {
+ loop {
+ let logs = harness.server().stdout_plain();
+ let mut counts = vec![0; shard_count];
+ let mut installed = HashSet::new();
+ for line in logs.lines().filter(|line|
line.contains(&install_marker)) {
+ let fields: Vec<_> = line.split_whitespace().collect();
+ let full_id: u128 = fields
+ .iter()
+ .find_map(|field| field.strip_prefix("client_id="))
+ .expect("install log must carry client_id")
+ .parse()
+ .unwrap();
+ let wire_id = u32::try_from(full_id &
u128::from(u32::MAX)).unwrap();
+ if !client_ids.contains(&wire_id) {
+ continue;
+ }
+ let owner = usize::try_from(full_id >> 112).unwrap();
+ assert!(owner < shard_count, "invalid owner: {line}");
+ assert!(
+ fields.contains(&format!("shard={owner}").as_str()),
+ "{line}"
+ );
+ assert!(
+ fields.contains(&format!("shard-{owner}").as_str()),
+ "{line}"
+ );
+ assert!(installed.insert(wire_id), "client installed twice:
{line}");
+ counts[owner] += 1;
+ }
+ if installed == client_ids {
+ assert_eq!(counts, vec![clients.len() / shard_count;
shard_count]);
+ break;
+ }
+ tokio::time::sleep(TLS_POLL_INTERVAL).await;
+ }
+ })
+ .await
+ .expect("every encrypted client must be installed on its owning shard
thread");
+ client_ids
+}
+
+async fn stalled_tls_clients(harness: &TestHarness, shard_count: usize) ->
Vec<TcpStream> {
+ let address = match harness.transport().unwrap() {
+ TransportProtocol::Tcp => harness.server().tcp_addr().unwrap(),
+ TransportProtocol::WebSocket =>
harness.server().websocket_addr().unwrap(),
+ transport => panic!("expected TCP-TLS or WSS, got {transport:?}"),
+ };
+ let mut clients = Vec::with_capacity(shard_count);
+ for _ in 0..shard_count {
+ clients.push(TcpStream::connect(address).await.unwrap());
+ }
+ for _ in 0..shard_count {
+ let mut invalid = TcpStream::connect(address).await.unwrap();
+ invalid.write_all(b"GET / HTTP/1.1\r\n\r\n").await.unwrap();
+ }
+ clients
+}
diff --git a/core/message_bus/src/client_listener/mod.rs
b/core/message_bus/src/client_listener/mod.rs
index f0f739496..5e1bf80b5 100644
--- a/core/message_bus/src/client_listener/mod.rs
+++ b/core/message_bus/src/client_listener/mod.rs
@@ -24,12 +24,10 @@
//! work runs in the install path (per [`crate::installer`]) so a slow
//! peer cannot block the accept loop.
//! - The `on_accepted` callback is responsible for minting a
-//! `client_id`. The plain TCP and pre-upgrade WS paths route to the
-//! owning shard via `shard::LifecycleFrame::ClientConnectionSetup` and
-//! `ClientWsConnectionSetup` respectively. TCP-TLS, WSS, and QUIC
-//! stay shard-0 terminal: their connection state is not serialisable
-//! so the callback installs locally on shard 0 instead of shipping a
-//! setup frame.
+//! `client_id`. TCP, WS, TCP-TLS and WSS delegate the raw TCP fd to the
+//! owning shard before any handshake. TLS configuration accompanies
+//! TCP-TLS and WSS setup frames. QUIC installs locally on shard 0
+//! because its connections share the endpoint's UDP socket.
//!
//! `bind` signature families:
//!
@@ -102,7 +100,7 @@ pub mod ws;
pub mod wss;
/// Bind a TCP listener with `TCP_NODELAY` set, the shared shape used by
-/// the plain-TCP and WS pre-upgrade client listeners.
+/// the plain-TCP, TCP-TLS and WS pre-upgrade client listeners.
///
/// Binding stays synchronous through `bind_reusable_tcp_listener` so shard
/// startup does not depend on compio's `IORING_OP_BIND` and
@@ -141,7 +139,8 @@ mod tests {
use server_common::executor::create_shard_executor;
use socket2::SockRef;
- use super::bind_nodelay_listener;
+ use super::{tcp, tcp_tls};
+ use crate::transports::tls::self_signed_for_loopback;
#[test]
fn
given_a_shard_executor_when_binding_a_client_listener_should_preserve_nodelay()
{
@@ -149,10 +148,20 @@ mod tests {
runtime.block_on(async {
let addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 0);
- let (listener, bound_addr) = bind_nodelay_listener(addr).unwrap();
+ let (tcp_listener, tcp_addr) = tcp::bind(addr).unwrap();
+ let (tls_listener, _, tls_addr) =
+ tcp_tls::bind(addr, self_signed_for_loopback()).unwrap();
- assert_ne!(bound_addr.port(), 0);
- assert!(SockRef::from(&listener).tcp_nodelay().unwrap());
+ for (transport, listener, bound_addr) in [
+ ("TCP", tcp_listener, tcp_addr),
+ ("TCP-TLS", tls_listener, tls_addr),
+ ] {
+ assert_ne!(bound_addr.port(), 0);
+ assert!(
+ SockRef::from(&listener).tcp_nodelay().unwrap(),
+ "{transport} listener must disable Nagle"
+ );
+ }
});
}
}
diff --git a/core/message_bus/src/client_listener/tcp_tls.rs
b/core/message_bus/src/client_listener/tcp_tls.rs
index 28ed0248d..cb3007633 100644
--- a/core/message_bus/src/client_listener/tcp_tls.rs
+++ b/core/message_bus/src/client_listener/tcp_tls.rs
@@ -24,13 +24,9 @@
//! ([`crate::installer::install_client_tcp_tls`]) so a slow or
//! malicious peer cannot block subsequent accepts.
//!
-//! TCP-TLS is shard-0 terminal. Pre-handshake the fd is plain TCP and
-//! could in principle be dup'd to another shard, but the receiving shard
-//! would then have to perform the rustls handshake against a config that
-//! lives on shard 0 — losing the point of the cross-shard handover.
-//! Post-handshake the rustls connection state holds per-record sequence
-//! numbers, key schedule, and write buffers tied to the local task, with
-//! no dupable plaintext fd.
+//! The coordinator delegates the raw fd and shared configuration before
+//! any TLS bytes are consumed. The destination shard creates all
+//! per-connection TLS state and owns subsequent encrypted I/O.
//!
//! The TLS plane structurally cannot preserve `Frozen<MESSAGE_ALIGN>`
//! ownership: rustls's encrypt step copies plaintext bytes into the
@@ -38,8 +34,8 @@
//! TCP does not carry over.
use crate::AcceptedTlsClientFn;
+use crate::client_listener::bind_nodelay_listener;
use crate::lifecycle::ShutdownToken;
-use crate::socket_opts::bind_reusable_tcp_listener;
use crate::transports::tls::{TlsServerCredentials,
install_default_crypto_provider};
use compio::net::TcpListener;
use futures::FutureExt;
@@ -67,8 +63,8 @@ use tracing::{debug, error, info};
///
/// # Errors
///
-/// - [`IggyError::IoError`] if the rustls server config cannot be built
-/// from `credentials` (cert / key mismatch).
+/// - [`IggyError::IoError`] if building the TLS config, configuring
+/// `TCP_NODELAY`, or reading the bound address fails.
/// - [`IggyError::CannotBindToSocket`] if the TCP bind fails.
#[allow(clippy::future_not_send)]
pub fn bind(
@@ -86,11 +82,7 @@ pub fn bind(
// cannot enable it accidentally.
cfg.max_early_data_size = 0;
- let listener = bind_reusable_tcp_listener(addr)
- .map_err(|_| IggyError::CannotBindToSocket(addr.to_string()))?;
- let actual = listener
- .local_addr()
- .map_err(|e| IggyError::IoError(e.to_string()))?;
+ let (listener, actual) = bind_nodelay_listener(addr)?;
Ok((listener, Arc::new(cfg), actual))
}
@@ -99,9 +91,8 @@ pub fn bind(
/// Each accepted [`compio::net::TcpStream`] is handed to `on_accepted`
/// together with a clone of the shared [`Arc<rustls::ServerConfig>`].
/// The callback owns the stream from that point on; production wiring
-/// routes through shard 0's coordinator (mints a `client_id`, builds
-/// the install context, calls
-/// [`crate::installer::install_client_tcp_tls`]).
+/// delegates the raw fd and configuration through shard 0's coordinator.
+/// The destination shard calls [`crate::installer::install_client_tcp_tls`].
#[allow(clippy::future_not_send)]
pub async fn run(
listener: TcpListener,
diff --git a/core/message_bus/src/client_listener/wss.rs
b/core/message_bus/src/client_listener/wss.rs
index f686c65f8..8308a4460 100644
--- a/core/message_bus/src/client_listener/wss.rs
+++ b/core/message_bus/src/client_listener/wss.rs
@@ -26,14 +26,9 @@
//! [`crate::transports::wss::WssTransportConn`]'s `run` method) so a
//! slow or malicious peer cannot block subsequent accepts.
//!
-//! WSS is shard-0 terminal. The plain TCP fd that exists pre-handshake
-//! could in principle be dup'd to another shard, but the receiving
-//! shard would then have to perform both the TLS handshake and the WS
-//! upgrade against state that lives on shard 0 — losing the point of
-//! the cross-shard handover. Post-upgrade the rustls + tungstenite
-//! state machines hold per-record sequence numbers, key schedule,
-//! masking state, and write buffers tied to the local task, with no
-//! dupable plaintext fd.
+//! The coordinator delegates the raw fd and shared configuration before
+//! either handshake. The destination shard creates all TLS and WebSocket
+//! state and owns subsequent encrypted I/O.
//!
//! Like the TCP-TLS plane, WSS structurally cannot preserve
//! `Frozen<MESSAGE_ALIGN>` ownership: rustls's encrypt step copies
@@ -90,9 +85,8 @@ pub fn bind(
/// Each accepted [`compio::net::TcpStream`] is handed to `on_accepted`
/// together with a clone of the shared [`Arc<rustls::ServerConfig>`].
/// The callback owns the stream from that point on; production wiring
-/// routes through shard 0's coordinator (mints a `client_id`, builds
-/// the install context, calls
-/// [`crate::installer::install_client_wss`]).
+/// delegates the raw fd and configuration through shard 0's coordinator.
+/// The destination shard calls [`crate::installer::install_client_wss`].
#[allow(clippy::future_not_send)]
pub async fn run(
listener: TcpListener,
diff --git a/core/message_bus/src/fd_transfer.rs
b/core/message_bus/src/fd_transfer.rs
index 118a5c56e..8bfd1f205 100644
--- a/core/message_bus/src/fd_transfer.rs
+++ b/core/message_bus/src/fd_transfer.rs
@@ -17,8 +17,8 @@
//! File descriptor transfer between shards.
//!
-//! After shard 0 accepts or connects a TCP socket and completes the
-//! handshake, it calls [`dup_fd`] to create a second kernel reference
+//! After shard 0 accepts or connects a TCP socket, before any handshake,
+//! it calls [`dup_fd`] to create a second kernel reference
//! to the same socket. The duplicated fd is wrapped in an owning
//! [`DupedFd`] and sent to the target shard via the inter-shard channel.
//! The target shard calls [`wrap_duped_fd`] to construct a compio
@@ -138,7 +138,9 @@ fn close_fd(fd: RawFd) {
#[cfg(test)]
mod tests {
use super::*;
+ use compio::io::AsyncRead;
use compio::net::{TcpListener, TcpStream};
+ use std::time::Duration;
#[compio::test]
#[allow(clippy::future_not_send)]
@@ -167,26 +169,28 @@ mod tests {
#[compio::test]
#[allow(clippy::future_not_send)]
async fn duped_fd_drops_close_underlying_fd() {
+ const CLOSE_TIMEOUT: Duration = Duration::from_secs(2);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let connect = TcpStream::connect(addr);
let accept = listener.accept();
let (client_res, accept_res) = futures::join!(connect, accept);
- let (_server, _) = accept_res.unwrap();
+ let (mut server, _) = accept_res.unwrap();
let client = client_res.unwrap();
let duped = dup_fd(&client).expect("dup_fd failed");
- let raw = duped.as_raw_fd();
+ drop(client);
drop(duped);
- // After drop the fd must be gone from this process' fd table.
- // SAFETY: F_GETFD on a closed fd is defined and returns -1/EBADF.
- let flags = unsafe { libc::fcntl(raw, libc::F_GETFD) };
- assert_eq!(flags, -1, "fd must be closed after DupedFd drop");
+ // Closed fd numbers can be reused immediately by other test threads.
+ let result = compio::time::timeout(CLOSE_TIMEOUT, server.read(vec![0]))
+ .await
+ .expect("DupedFd drop must release the socket")
+ .0;
assert_eq!(
- io::Error::last_os_error().raw_os_error(),
- Some(libc::EBADF),
- "closed fd must report EBADF"
+ result.unwrap(),
+ 0,
+ "peer must observe EOF after both socket handles close"
);
}
}
diff --git a/core/message_bus/src/installer/common.rs
b/core/message_bus/src/installer/common.rs
index 918c23ad5..9d0278b75 100644
--- a/core/message_bus/src/installer/common.rs
+++ b/core/message_bus/src/installer/common.rs
@@ -28,11 +28,8 @@ use std::time::{Duration, Instant};
/// already told both tasks to skip post-loop cleanup. Closing the
/// sender wakes the transport's writer; triggering the per-connection
/// shutdown wakes the transport's reader off its `io_uring` read SQE
-/// without waiting for peer EOF. Awaiting both handles with
-/// `close_peer_timeout` budget guarantees neither task can outlive the
-/// race on a half-open socket. `compio::runtime::JoinHandle::drop`
-/// detaches, so letting the handles go out of scope would leak the
-/// tasks.
+/// without waiting for peer EOF. Awaiting both handles allows cooperative
+/// cleanup; dropping a timed-out handle cancels its task.
///
/// Both handles share a single `timeout` budget: after the transport
/// returns (or is cancelled) the dispatch only gets the remaining time.
diff --git a/core/message_bus/src/installer/mod.rs
b/core/message_bus/src/installer/mod.rs
index 641fd74e6..7ea17868b 100644
--- a/core/message_bus/src/installer/mod.rs
+++ b/core/message_bus/src/installer/mod.rs
@@ -50,16 +50,16 @@ use crate::client_listener::RequestHandler;
use crate::fd_transfer::{self, DupedFd};
use crate::installer::conn_info::ClientConnMeta;
use crate::replica::listener::MessageHandler;
-use crate::{ClientConnectionLostFn, IggyMessageBus, ReplicaHandshakeDoneFn};
+use crate::{
+ ClientConnectionLostFn, IggyMessageBus, ReplicaHandshakeDoneFn,
SharedTlsServerConfig,
+};
use std::rc::Rc;
use tracing::warn;
/// Operations a shard needs to perform on its local bus when the router
/// receives an inter-shard connection-setup frame.
///
-/// The production implementation is on `Rc<IggyMessageBus>`. The simulator
-/// does not exercise this path; if it ever does, add a no-op impl on
-/// `SharedSimOutbox`.
+/// The production implementation is on `Rc<IggyMessageBus>`.
pub trait ConnectionInstaller {
/// Wrap a blind-delegated inbound replica fd into a `TcpStream` on
/// the local compio runtime, run the acceptor handshake in a spawned
@@ -113,6 +113,26 @@ pub trait ConnectionInstaller {
/// access via the LOGIN allowlist.
fn install_client_ws_fd(&self, fd: DupedFd, meta: ClientConnMeta,
on_request: RequestHandler);
+ /// Wrap an unhandshaken TCP-TLS fd on this runtime and install its
+ /// transport task. The task owns the handshake, I/O and cleanup.
+ fn install_client_tcp_tls_fd(
+ &self,
+ fd: DupedFd,
+ meta: ClientConnMeta,
+ config: SharedTlsServerConfig,
+ on_request: RequestHandler,
+ );
+
+ /// Install an unhandshaken WSS fd on this runtime. TLS and the WebSocket
+ /// upgrade share the existing `handshake_grace` deadline in the task.
+ fn install_client_wss_fd(
+ &self,
+ fd: DupedFd,
+ meta: ClientConnMeta,
+ config: SharedTlsServerConfig,
+ on_request: RequestHandler,
+ );
+
/// Per-connection metadata stored at install time, or `None` if the client
/// is not (or no longer) connected on this bus. Dispatch reads `peer_addr`
/// + `transport` from it to seed the shard's session manager.
@@ -173,11 +193,7 @@ impl ConnectionInstaller for Rc<IggyMessageBus> {
)
.await;
match outcome {
- Ok(Ok(ws)) => {
- if !bus.is_shutting_down() {
- install_client_ws(&bus, meta, ws, on_request);
- }
- }
+ Ok(Ok(ws)) => install_client_ws(&bus, meta, ws, on_request),
Ok(Err(e)) => {
warn!(client_id = meta.client_id, "WS upgrade failed:
{e}");
}
@@ -193,6 +209,28 @@ impl ConnectionInstaller for Rc<IggyMessageBus> {
self.track_background(handle);
}
+ fn install_client_tcp_tls_fd(
+ &self,
+ fd: DupedFd,
+ meta: ClientConnMeta,
+ config: SharedTlsServerConfig,
+ on_request: RequestHandler,
+ ) {
+ let stream = fd_transfer::wrap_duped_fd(fd);
+ install_client_tcp_tls(self, meta, stream, config, on_request);
+ }
+
+ fn install_client_wss_fd(
+ &self,
+ fd: DupedFd,
+ meta: ClientConnMeta,
+ config: SharedTlsServerConfig,
+ on_request: RequestHandler,
+ ) {
+ let stream = fd_transfer::wrap_duped_fd(fd);
+ install_client_wss(self, meta, stream, config, on_request);
+ }
+
fn client_meta(&self, client_id: u128) -> Option<Rc<ClientConnMeta>> {
IggyMessageBus::client_meta(self, client_id)
}
diff --git a/core/message_bus/src/installer/replica.rs
b/core/message_bus/src/installer/replica.rs
index bf8a344b7..5bb4556fb 100644
--- a/core/message_bus/src/installer/replica.rs
+++ b/core/message_bus/src/installer/replica.rs
@@ -369,8 +369,6 @@ pub fn install_replica_conn<C: TransportConn>(
// `notify_connection_lost` stands down whenever a live registry entry
// exists - but the loser also drives `replica_dispatch_loop`, which
// must never hand the winner's replica id to `on_message`.
- // `compio::runtime::JoinHandle::drop` does not cancel the spawned
- // task, so we have to tell the tasks to stand down in-band.
let install_aborted = Rc::new(Cell::new(false));
// Generation token published by the registry on a successful insert.
@@ -575,15 +573,9 @@ pub fn install_replica_conn<C: TransportConn>(
install_aborted.set(true);
let _ = bus.clear_replica_owned(peer_id);
warn!(replica = peer_id, "replica registry insert raced");
- // `drain_rejected_registration` triggers the per-connection
- // shutdown (returned with `rejected`) to wake the transport
- // off its `io_uring` read SQE, then awaits transport /
- // dispatch handles. `compio::runtime::JoinHandle::drop` only
- // detaches, so without this drain the reader would outlive
- // the race on a half-open socket until peer EOF. Hand the
- // handle to `track_background` so `IggyMessageBus::shutdown`
- // awaits the drain before returning - `.detach()` would
- // orphan it and leak the half-closed socket across shutdown.
+ // Allow bounded cooperative cleanup of the losing connection.
+ // Track the drain so bus shutdown awaits it; dropping its
+ // handle would cancel that cleanup.
let drain_handle =
compio::runtime::spawn(drain_rejected_registration(rejected,
close_peer_timeout));
bus.track_background(drain_handle);
diff --git a/core/message_bus/src/installer/tcp.rs
b/core/message_bus/src/installer/tcp.rs
index 67bfc091e..b9908bae8 100644
--- a/core/message_bus/src/installer/tcp.rs
+++ b/core/message_bus/src/installer/tcp.rs
@@ -61,6 +61,8 @@ pub fn install_client_tcp(
/// Install a pre-wrapped client connection on the bus. Generic over
/// [`TransportConn`]; plane-symmetric with
/// [`super::replica::install_replica_conn`].
+///
+/// Connections arriving after shutdown starts are dropped before registration.
#[allow(clippy::future_not_send, clippy::too_many_lines)]
pub fn install_client_conn<C: TransportConn>(
bus: &Rc<IggyMessageBus>,
@@ -68,6 +70,9 @@ pub fn install_client_conn<C: TransportConn>(
conn: C,
on_request: RequestHandler,
) {
+ if bus.is_shutting_down() {
+ return;
+ }
let client_id = meta.client_id;
let (tx, rx) = async_channel::bounded(bus.config().client_queue_capacity);
let (in_tx, in_rx) =
diff --git a/core/message_bus/src/installer/tcp_tls.rs
b/core/message_bus/src/installer/tcp_tls.rs
index 3497aac79..5db4c84df 100644
--- a/core/message_bus/src/installer/tcp_tls.rs
+++ b/core/message_bus/src/installer/tcp_tls.rs
@@ -43,12 +43,9 @@ use tracing::warn;
/// to accepted sockets, so toggling here is required. `SO_KEEPALIVE`
/// is intentionally NOT set; see `socket_opts`.
///
-/// TCP-TLS is shard-0 terminal: the rustls connection state machine
-/// is non-serialisable and tied to the local task; pre-handshake the
-/// fd is plain TCP and could in principle be dup'd to another shard,
-/// but the receiving shard would then have to re-handshake against
-/// shard-0-resident key material — losing the point of the cross-shard
-/// handover.
+/// Called on the destination shard after raw-fd delegation. The shared
+/// configuration crosses shards; per-connection TLS state stays on this
+/// runtime for the lifetime of the connection.
#[allow(clippy::future_not_send)]
pub fn install_client_tcp_tls(
bus: &Rc<IggyMessageBus>,
diff --git a/core/message_bus/src/installer/wss.rs
b/core/message_bus/src/installer/wss.rs
index d65f66fd3..25f6ee5d9 100644
--- a/core/message_bus/src/installer/wss.rs
+++ b/core/message_bus/src/installer/wss.rs
@@ -43,9 +43,8 @@ use tracing::warn;
/// [`super::tcp_tls::install_client_tcp_tls`]. `SO_KEEPALIVE` is
/// intentionally NOT set; see `socket_opts`.
///
-/// WSS is shard-0 terminal for the same reasons as the TCP-TLS plane;
-/// see [`super::tcp_tls::install_client_tcp_tls`] for the rustls
-/// non-serialisability argument.
+/// Called on the destination shard after raw-fd delegation. Both TLS and
+/// WebSocket state are created and driven exclusively on this runtime.
#[allow(clippy::future_not_send)]
pub fn install_client_wss(
bus: &Rc<IggyMessageBus>,
diff --git a/core/message_bus/src/lib.rs b/core/message_bus/src/lib.rs
index d8038b612..dc2348daa 100644
--- a/core/message_bus/src/lib.rs
+++ b/core/message_bus/src/lib.rs
@@ -57,10 +57,10 @@
//! (their `Frozen<MESSAGE_ALIGN>` fragments flattened, chunked at
//! `IOV_MAX`) into `write_vectored_all`. Don't introduce per-message
//! syscalls or per-message encryption on the plaintext TCP plane.
-//! - fd-delegation ([`fd_transfer`]) is TCP-only. TLS / QUIC
-//! connections have no dupable plaintext fd, so shard 0 terminates
-//! and forwards `Frozen<MESSAGE_ALIGN>` over the existing
-//! inter-shard crossfire channel.
+//! - fd-delegation ([`fd_transfer`]) transfers raw TCP sockets for TCP,
+//! WS, TCP-TLS and WSS before handshakes. The destination shard owns
+//! handshake state and I/O. QUIC client callbacks install on shard 0
+//! through its shared UDP endpoint.
//! - 0-RTT stays disabled by default on any future QUIC path. Per-
//! command opt-in requires a checked-in idempotence audit.
//!
@@ -403,42 +403,26 @@ pub type AcceptedQuicClientFn = std::rc::Rc<dyn
Fn(AcceptedQuicConn)>;
/// requires a dupable plaintext fd) stays well-defined.
pub type AcceptedWsClientFn = std::rc::Rc<dyn Fn(compio::net::TcpStream)>;
+/// Listener TLS configuration shared with the shard owning each connection.
+/// Per-connection TLS state is created only on the destination runtime.
+pub type SharedTlsServerConfig = std::sync::Arc<rustls::ServerConfig>;
+
/// Callback invoked on every accepted SDK TCP-TLS client connection.
///
/// Fires after shard 0's TCP-TLS listener accepts a raw TCP socket.
-/// Neither the rustls handshake nor any application-layer work has run
-/// yet — the listener stays cheap so a slow handshake on one peer cannot
-/// block subsequent accepts. The callback receives the raw stream plus
-/// a clone of the shared [`std::sync::Arc<rustls::ServerConfig>`] built
-/// at bind time, mints a `client_id`, and calls
-/// [`installer::install_client_tcp_tls`]; the install path drives the
-/// rustls handshake on its own task before forwarding the connection to
-/// [`installer::install_client_conn`].
-///
-/// TCP-TLS stays shard-0 terminal: rustls's connection state machine
-/// is tied to the local task and not serialisable, and the
-/// pre-handshake fd would have to re-handshake on the receiving shard,
-/// losing the point of fd-delegation.
-pub type AcceptedTlsClientFn =
- std::rc::Rc<dyn Fn(compio::net::TcpStream,
std::sync::Arc<rustls::ServerConfig>)>;
+/// The callback delegates the socket and its listener's configuration
+/// before any TLS bytes are consumed. The destination shard runs
+/// [`installer::install_client_tcp_tls`], owning the handshake and all I/O.
+pub type AcceptedTlsClientFn = std::rc::Rc<dyn Fn(compio::net::TcpStream,
SharedTlsServerConfig)>;
/// Callback invoked on every accepted SDK WSS (WebSocket-over-TLS)
/// client connection.
///
-/// Fires after shard 0's WSS listener accepts a raw TCP socket. Neither
-/// the rustls handshake nor the WebSocket HTTP-Upgrade has run yet — the
-/// listener stays cheap so neither handshake on one peer can block
-/// subsequent accepts. The callback receives the raw stream plus a clone
-/// of the shared [`std::sync::Arc<rustls::ServerConfig>`] built at bind
-/// time, mints a `client_id`, and calls
-/// [`installer::install_client_wss`]; the install path drives
-/// both handshakes on its own task before forwarding the connection to
-/// [`installer::install_client_conn`].
-///
-/// No subprotocol negotiation is performed. WSS stays shard-0 terminal
-/// for the same reasons as the TCP-TLS plane.
-pub type AcceptedWssClientFn =
- std::rc::Rc<dyn Fn(compio::net::TcpStream,
std::sync::Arc<rustls::ServerConfig>)>;
+/// Fires before either the TLS handshake or WebSocket upgrade. The callback
+/// delegates the raw socket and its listener's configuration to the owning
+/// shard, where [`installer::install_client_wss`] runs both handshakes and
+/// subsequent I/O. No subprotocol negotiation is performed.
+pub type AcceptedWssClientFn = std::rc::Rc<dyn Fn(compio::net::TcpStream,
SharedTlsServerConfig)>;
/// Notifier fired when a delegated replica connection dies.
///
@@ -1224,6 +1208,8 @@ impl IggyMessageBus {
let remaining =
deadline.saturating_duration_since(std::time::Instant::now());
let clients_outcome = self.clients.drain(remaining).await;
+ // Connection tasks skip per-client cleanup during bus shutdown.
+ self.client_meta.borrow_mut().clear();
let remaining =
deadline.saturating_duration_since(std::time::Instant::now());
let replicas_outcome = self.replicas.drain(remaining).await;
diff --git a/core/message_bus/src/lifecycle/connection_registry.rs
b/core/message_bus/src/lifecycle/connection_registry.rs
index 24c709742..571e615eb 100644
--- a/core/message_bus/src/lifecycle/connection_registry.rs
+++ b/core/message_bus/src/lifecycle/connection_registry.rs
@@ -90,12 +90,11 @@ pub type BusReceiver = async_channel::Receiver<BusMessage>;
///
/// Exposes the writer and reader [`JoinHandle`]s plus the per-connection
/// [`Shutdown`] so the loser can explicitly drain (or force-cancel on
-/// deadline) the orphan tasks rather than relying on them to self-exit via
-/// `install_aborted`. `compio::runtime::JoinHandle::drop` detaches; without
-/// this the loser would leak the handles and a reader looping on
-/// `framing::read_message` could outlive the race indefinitely on a
-/// half-open socket. Triggering the [`Shutdown`] wakes the reader off its
-/// `io_uring` read SQE without waiting for peer EOF.
+/// deadline) the orphan tasks. `install_aborted` prevents them from
+/// dispatching frames or cleaning up the winning entry. Dropping a
+/// [`JoinHandle`] cancels its task; explicitly draining allows cooperative
+/// shutdown before the deadline. Triggering [`Shutdown`] wakes the reader
+/// without waiting for peer EOF.
///
/// [`Shutdown`]: crate::lifecycle::Shutdown
#[derive(Debug)]
diff --git a/core/message_bus/src/replica/io.rs
b/core/message_bus/src/replica/io.rs
index 9368ef5a7..0ae9e6874 100644
--- a/core/message_bus/src/replica/io.rs
+++ b/core/message_bus/src/replica/io.rs
@@ -18,25 +18,20 @@
//! Shard-0 listener bootstrap.
//!
//! Every shard on a node instantiates its own `Rc<IggyMessageBus>` on its
-//! own compio runtime, but only shard 0 terminates inbound traffic: it
-//! binds the replica listener, the plain-TCP client listener, and
-//! optionally the WS, QUIC, TCP-TLS, and WSS client listeners. It also
-//! dials higher-id peers for the replica plane. Each accepted or
+//! own compio runtime. Shard 0 binds the replica and client listeners
+//! and dials higher-id peers for the replica plane. Each accepted or
//! dialed connection is handed to the delegate callback supplied by
//! the caller (typically wrapping `shard::coordinator::ShardZeroCoordinator`).
-//! Replica, client-TCP, and pre-upgrade WS callbacks duplicate the raw
-//! fd and ship it to the owning shard via the inter-shard `ShardFrame`
-//! channel; replica delegation is blind (no byte is read on shard 0 -
-//! the owning shard runs the handshake, see
-//! [`crate::replica::handshake`]). TCP-TLS, WSS, and QUIC CLIENT
-//! callbacks install locally on shard 0: their handshake state
-//! materialises before delegation could happen, and post-handshake
-//! rustls / QUIC state is not serialisable.
+//! Replica, TCP, WS, TCP-TLS and WSS callbacks duplicate the raw fd and
+//! ship it to the owning shard before any bytes are read. TLS configuration
+//! travels with TCP-TLS and WSS sockets. The owning shard runs handshakes
+//! and subsequent I/O; see [`crate::replica::handshake`] for replica
+//! handshakes. QUIC client callbacks install locally on shard 0 through
+//! its shared UDP endpoint.
//!
-//! Non-zero shards early-return `Ok(None)`; the launcher calls this helper
-//! unconditionally per shard, so non-zero shards just have no listener
-//! binding and rely on `send_to_*` slow-path forwarding to reach the
-//! owning shard.
+//! Non-zero shards early-return `Ok(None)` and receive delegated sockets
+//! through their inboxes. `send_to_*` forwards traffic when another shard
+//! owns the connection.
use std::net::SocketAddr;
use std::rc::Rc;
@@ -135,17 +130,13 @@ pub fn assert_listen_addrs_distinct(
/// Optionally bind WS, QUIC, TCP-TLS, and WSS client listeners
/// alongside. Non-zero shards early-return `Ok(None)`.
///
-/// Each accepted / dialed connection is handed to the supplied
-/// delegate callback. The replica callbacks (`on_accepted_replica` for
-/// blind inbound delegation, `on_dialed_replica` for outbound), the
-/// client-TCP callback (`on_accepted_client`), and the WS callback
-/// (`on_accepted_ws_client`) are responsible for the dup-fd +
-/// inter-shard send. The QUIC, TCP-TLS, and WSS callbacks
-/// (`on_accepted_quic_client` / `on_accepted_tcp_tls_client` /
-/// `on_accepted_wss_client`) install locally on shard 0; those CLIENT
-/// planes have no cross-shard handover (QUIC by transport design,
-/// TLS-family because the handshake completes before delegation could
-/// happen and the resulting rustls state is not serialisable).
+/// Each accepted or dialed connection is handed to its supplied callback.
+/// Replica, TCP, WS, TCP-TLS and WSS callbacks duplicate the raw fd and
+/// send it to the owning shard before any handshake. TCP-TLS and WSS
+/// also transfer the listener's shared TLS configuration. The owning
+/// shard supplies its handlers and performs handshakes and I/O.
+/// The QUIC client callback installs locally on shard 0 through its
+/// shared UDP endpoint.
///
/// `ws_listen_addr` / `on_accepted_ws_client` are paired: if either is
/// `Some`, both must be. Same for the QUIC trio
diff --git a/core/message_bus/src/transports/tcp_tls.rs
b/core/message_bus/src/transports/tcp_tls.rs
index c3fa5ac21..c0155991a 100644
--- a/core/message_bus/src/transports/tcp_tls.rs
+++ b/core/message_bus/src/transports/tcp_tls.rs
@@ -218,7 +218,11 @@ impl TransportConn for TcpTlsTransportConn {
let mut tls = match self.state {
ConnState::Established(tls) => *tls,
ConnState::Pending { stream, role } => {
- match compio::time::timeout(handshake_grace, handshake(role,
stream)).await {
+ let outcome = futures::select_biased! {
+ () = ctx.shutdown.wait().fuse() => return,
+ outcome = compio::time::timeout(handshake_grace,
handshake(role, stream)).fuse() => outcome,
+ };
+ match outcome {
Ok(Ok(s)) => s,
Ok(Err(e)) => {
warn!(
diff --git a/core/message_bus/src/transports/wss.rs
b/core/message_bus/src/transports/wss.rs
index 20168d008..b28b93ebf 100644
--- a/core/message_bus/src/transports/wss.rs
+++ b/core/message_bus/src/transports/wss.rs
@@ -174,12 +174,14 @@ impl TransportConn for WssTransportConn {
// both inner handshake futures are large (rustls + tungstenite
// state machines) and would otherwise push `run`'s state past
// the `clippy::large_futures` threshold.
- let tls_stream = match compio::time::timeout(
- handshake_grace,
- Box::pin(tls_handshake(role, self.stream)),
- )
- .await
- {
+ let tls_outcome = futures::select_biased! {
+ () = ctx.shutdown.wait().fuse() => return,
+ outcome = compio::time::timeout(
+ handshake_grace,
+ Box::pin(tls_handshake(role, self.stream)),
+ ).fuse() => outcome,
+ };
+ let tls_stream = match tls_outcome {
Ok(Ok(s)) => s,
Ok(Err(e)) => {
warn!(
@@ -212,12 +214,14 @@ impl TransportConn for WssTransportConn {
);
return;
}
- let mut ws = match compio::time::timeout(
- remaining,
- Box::pin(ws_handshake(tls_stream, is_server, &peer, ws_config)),
- )
- .await
- {
+ let upgrade_outcome = futures::select_biased! {
+ () = ctx.shutdown.wait().fuse() => return,
+ outcome = compio::time::timeout(
+ remaining,
+ Box::pin(ws_handshake(tls_stream, is_server, &peer,
ws_config)),
+ ).fuse() => outcome,
+ };
+ let mut ws = match upgrade_outcome {
Ok(Ok(s)) => s,
Ok(Err(e)) => {
warn!(
diff --git a/core/message_bus/tests/common/mod.rs
b/core/message_bus/tests/common/mod.rs
index b5630b764..da16e8dd9 100644
--- a/core/message_bus/tests/common/mod.rs
+++ b/core/message_bus/tests/common/mod.rs
@@ -26,6 +26,8 @@
#![allow(dead_code)] // each test binary uses a subset
+use compio::io::{AsyncRead, AsyncWriteExt};
+use compio::net::TcpStream;
use iggy_binary_protocol::{Command, GenericHeader, HEADER_SIZE};
use message_bus::ConnectionInstaller;
use message_bus::client_listener::RequestHandler;
@@ -46,6 +48,10 @@ use std::cell::Cell;
use std::net::SocketAddr;
use std::rc::Rc;
use std::sync::Arc;
+use std::time::Duration;
+
+const CLIENT_CLEANUP_TIMEOUT: Duration = Duration::from_secs(2);
+const CLIENT_POLL_INTERVAL: Duration = Duration::from_millis(10);
/// Build a stub [`ClientConnMeta`] for tests that don't care about
/// peer addr / transport details. Uses `127.0.0.1:0` and the given
@@ -263,9 +269,9 @@ pub fn install_ws_clients_locally(
}
/// Build an [`AcceptedTlsClientFn`] that mints a local client id and
-/// installs the accepted TCP-TLS stream directly on the given bus. The
-/// install path drives the rustls handshake on its own task, mirroring
-/// the production shard-0 coordinator.
+/// duplicates the raw TCP-TLS fd into the owning-shard installer. This
+/// exercises fd wrapping locally; server integration tests cover routing
+/// from the production accept callback to a different shard.
#[must_use]
pub fn install_tls_clients_locally(
bus: Rc<IggyMessageBus>,
@@ -278,12 +284,14 @@ pub fn install_tls_clients_locally(
counter.set(seq.wrapping_add(1));
let client_id = (shard_id << 112) | seq;
let meta = test_client_meta(client_id, ClientTransportKind::TcpTls);
- installer::install_client_tcp_tls(&bus, meta, stream, config,
on_request.clone());
+ let fd = fd_transfer::dup_fd(&stream).expect("dup_fd");
+ drop(stream);
+ bus.install_client_tcp_tls_fd(fd, meta, config, on_request.clone());
})
}
/// Build an [`AcceptedWssClientFn`] that mints a local client id and
-/// installs the accepted WSS stream directly on the given bus. The
+/// duplicates the raw WSS fd into the owning-shard installer. The
/// install path drives both the rustls handshake and the WS HTTP-Upgrade
/// inside the transport's `run` body.
#[must_use]
@@ -298,6 +306,97 @@ pub fn install_wss_clients_locally(
counter.set(seq.wrapping_add(1));
let client_id = (shard_id << 112) | seq;
let meta = test_client_meta(client_id, ClientTransportKind::Wss);
- installer::install_client_wss(&bus, meta, stream, config,
on_request.clone());
+ let fd = fd_transfer::dup_fd(&stream).expect("dup_fd");
+ drop(stream);
+ bus.install_client_wss_fd(fd, meta, config, on_request.clone());
})
}
+
+#[allow(clippy::future_not_send)]
+pub async fn wait_for_client_count(bus: &IggyMessageBus, count: usize) {
+ compio::time::timeout(CLIENT_CLEANUP_TIMEOUT, async {
+ while bus.clients().len() != count {
+ compio::time::sleep(CLIENT_POLL_INTERVAL).await;
+ }
+ })
+ .await
+ .unwrap_or_else(|_| panic!("expected {count} clients, got {}",
bus.clients().len()));
+}
+
+#[allow(clippy::future_not_send)]
+pub async fn assert_failed_tls_installs_are_removed(transport:
ClientTransportKind) {
+ const OWNER: u16 = 1;
+ const FIRST_CLIENT_ID: u128 = (OWNER as u128) << 112 | 1;
+ const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(1);
+ let bus = Rc::new(IggyMessageBus::new(OWNER));
+ let request_seen = Rc::new(Cell::new(false));
+ let seen_by_handler = Rc::clone(&request_seen);
+ let on_request: RequestHandler = Rc::new(move |_, _|
seen_by_handler.set(true));
+ let credentials = self_signed_for_loopback();
+ let (listener, config, address, accepted) = match transport {
+ ClientTransportKind::TcpTls => {
+ let (listener, config, address) =
+ message_bus::client_listener::tcp_tls::bind(loopback(),
credentials).unwrap();
+ (
+ listener,
+ config,
+ address,
+ install_tls_clients_locally(Rc::clone(&bus), on_request),
+ )
+ }
+ ClientTransportKind::Wss => {
+ let (listener, config, address) =
+ message_bus::client_listener::wss::bind(loopback(),
credentials).unwrap();
+ (
+ listener,
+ config,
+ address,
+ install_wss_clients_locally(Rc::clone(&bus), on_request),
+ )
+ }
+ _ => panic!("expected encrypted TCP transport"),
+ };
+
+ let mut invalid = TcpStream::connect(address).await.unwrap();
+ accepted(listener.accept().await.unwrap().0, Arc::clone(&config));
+ assert!(bus.client_meta(FIRST_CLIENT_ID).is_some());
+ invalid
+ .write_all(b"GET / HTTP/1.1\r\n\r\n")
+ .await
+ .0
+ .unwrap();
+ wait_for_client_count(&bus, 0).await;
+ assert!(bus.client_meta(FIRST_CLIENT_ID).is_none());
+
+ let stalled = TcpStream::connect(address).await.unwrap();
+ accepted(listener.accept().await.unwrap().0, Arc::clone(&config));
+ assert!(bus.client_meta(FIRST_CLIENT_ID + 1).is_some());
+ // Poll the spawned transport so shutdown interrupts a pending handshake.
+ compio::time::sleep(CLIENT_POLL_INTERVAL).await;
+ assert_eq!(
+ bus.shutdown(SHUTDOWN_TIMEOUT).await.force,
+ 0,
+ "stalled {transport:?} handshake must observe shutdown without
force-cancellation"
+ );
+ assert!(bus.clients().is_empty());
+ assert!(bus.client_meta(FIRST_CLIENT_ID + 1).is_none());
+
+ let late = TcpStream::connect(address).await.unwrap();
+ accepted(listener.accept().await.unwrap().0, config);
+ assert!(
+ bus.clients().is_empty(),
+ "setup after shutdown must be rejected"
+ );
+ assert!(bus.client_meta(FIRST_CLIENT_ID + 2).is_none());
+ for mut peer in [stalled, late] {
+ let result = compio::time::timeout(CLIENT_CLEANUP_TIMEOUT,
peer.read(vec![0]))
+ .await
+ .expect("shutdown must close pending and late handshake sockets")
+ .0;
+ assert_eq!(result.unwrap(), 0);
+ }
+ assert!(
+ !request_seen.get(),
+ "failed handshakes must not dispatch requests"
+ );
+}
diff --git a/core/message_bus/tests/tcp_tls_client_listener.rs
b/core/message_bus/tests/tcp_tls_client_listener.rs
index d477d1cd4..1a826e4be 100644
--- a/core/message_bus/tests/tcp_tls_client_listener.rs
+++ b/core/message_bus/tests/tcp_tls_client_listener.rs
@@ -182,3 +182,8 @@ async fn slow_tls_handshake_evicts_registry() {
"graceful shutdown should not force-cancel"
);
}
+
+#[compio::test]
+async fn failed_and_interrupted_tls_handshakes_release_delegated_connections()
{
+
common::assert_failed_tls_installs_are_removed(message_bus::ClientTransportKind::TcpTls).await;
+}
diff --git a/core/message_bus/tests/wss_client_listener.rs
b/core/message_bus/tests/wss_client_listener.rs
index f280d8ebf..b65ea5522 100644
--- a/core/message_bus/tests/wss_client_listener.rs
+++ b/core/message_bus/tests/wss_client_listener.rs
@@ -19,7 +19,9 @@ mod common;
use async_channel::bounded;
use common::{header_only, install_wss_clients_locally, loopback};
+use compio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
use compio::net::TcpStream;
+use compio::tls::TlsConnector;
use iggy_binary_protocol::Command;
use iggy_binary_protocol::GenericHeader;
use message_bus::BusMessage;
@@ -32,6 +34,7 @@ use message_bus::{FusedShutdown, IggyMessageBus, MessageBus,
MessageBusConfig, S
use rustls::RootCertStore;
use rustls::pki_types::ServerName;
use server_common::Message;
+use std::cell::Cell;
use std::rc::Rc;
use std::sync::Arc;
use std::time::{Duration, Instant};
@@ -170,3 +173,108 @@ async fn slow_handshake_evicts_registry() {
"graceful shutdown should not force-cancel"
);
}
+
+#[compio::test]
+async fn failed_and_interrupted_wss_handshakes_release_delegated_connections()
{
+
common::assert_failed_tls_installs_are_removed(message_bus::ClientTransportKind::Wss).await;
+}
+
+#[compio::test]
+async fn
stalled_rejected_and_interrupted_websocket_upgrades_release_delegated_connections()
{
+ const OWNER: u16 = 1;
+ const FIRST_CLIENT_ID: u128 = (OWNER as u128) << 112 | 1;
+ const HANDSHAKE_GRACE: Duration = Duration::from_secs(2);
+ const TLS_DELAY: Duration = Duration::from_millis(1_200);
+ const DEADLINE_MARGIN: Duration = Duration::from_millis(400);
+ const IO_TIMEOUT: Duration = Duration::from_secs(3);
+ const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(1);
+ const UPGRADE_WAIT: Duration = Duration::from_millis(10);
+
+ install_default_crypto_provider();
+ let bus = Rc::new(IggyMessageBus::with_tunables(
+ OWNER,
+ MessageBusConfig {
+ handshake_grace: HANDSHAKE_GRACE,
+ ..MessageBusConfig::default()
+ },
+ ));
+ let credentials = self_signed_for_loopback();
+ let mut roots = RootCertStore::empty();
+ for cert in &credentials.cert_chain {
+ roots.add(cert.clone()).unwrap();
+ }
+ let connector = TlsConnector::from(Arc::new(
+ rustls::ClientConfig::builder()
+ .with_root_certificates(roots)
+ .with_no_client_auth(),
+ ));
+ let (listener, config, address) = bind(loopback(), credentials).unwrap();
+ let request_seen = Rc::new(Cell::new(false));
+ let seen_by_handler = Rc::clone(&request_seen);
+ let on_request: RequestHandler = Rc::new(move |_, _|
seen_by_handler.set(true));
+ let accepted = install_wss_clients_locally(Rc::clone(&bus), on_request);
+
+ let tcp = TcpStream::connect(address).await.unwrap();
+ accepted(listener.accept().await.unwrap().0, Arc::clone(&config));
+ let installed_at = Instant::now();
+ compio::time::sleep(TLS_DELAY).await;
+ let mut tls = compio::time::timeout(IO_TIMEOUT,
connector.connect("localhost", tcp))
+ .await
+ .unwrap()
+ .unwrap();
+ common::wait_for_client_count(&bus, 0).await;
+ assert!(
+ installed_at.elapsed() < HANDSHAKE_GRACE + DEADLINE_MARGIN,
+ "TLS and WebSocket must share one deadline, elapsed {:?}",
+ installed_at.elapsed()
+ );
+ assert!(bus.client_meta(FIRST_CLIENT_ID).is_none());
+ let closed = compio::time::timeout(IO_TIMEOUT, tls.read(vec![0]))
+ .await
+ .unwrap()
+ .0;
+ assert!(closed.is_err() || closed.unwrap() == 0);
+
+ let tcp = TcpStream::connect(address).await.unwrap();
+ accepted(listener.accept().await.unwrap().0, Arc::clone(&config));
+ let mut tls = compio::time::timeout(IO_TIMEOUT,
connector.connect("localhost", tcp))
+ .await
+ .unwrap()
+ .unwrap();
+ assert!(bus.client_meta(FIRST_CLIENT_ID + 1).is_some());
+ tls.write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n")
+ .await
+ .0
+ .unwrap();
+ tls.flush().await.unwrap();
+ common::wait_for_client_count(&bus, 0).await;
+ assert!(bus.client_meta(FIRST_CLIENT_ID + 1).is_none());
+
+ let tcp = TcpStream::connect(address).await.unwrap();
+ accepted(listener.accept().await.unwrap().0, config);
+ let mut tls = compio::time::timeout(IO_TIMEOUT,
connector.connect("localhost", tcp))
+ .await
+ .unwrap()
+ .unwrap();
+ tls.write_all(b"GET / HTTP/1.1\r\n").await.0.unwrap();
+ tls.flush().await.unwrap();
+ // Let the server finish TLS and wait for the remaining upgrade headers.
+ compio::time::sleep(UPGRADE_WAIT).await;
+ assert!(bus.client_meta(FIRST_CLIENT_ID + 2).is_some());
+ assert_eq!(
+ bus.shutdown(SHUTDOWN_TIMEOUT).await.force,
+ 0,
+ "pending WebSocket upgrade must observe shutdown without
force-cancellation"
+ );
+ assert!(bus.clients().is_empty());
+ assert!(bus.client_meta(FIRST_CLIENT_ID + 2).is_none());
+ let closed = compio::time::timeout(IO_TIMEOUT, tls.read(vec![0]))
+ .await
+ .unwrap()
+ .0;
+ assert!(closed.is_err() || closed.unwrap() == 0);
+ assert!(
+ !request_seen.get(),
+ "failed upgrades must not dispatch requests"
+ );
+}
diff --git a/core/server/src/boot/listeners.rs
b/core/server/src/boot/listeners.rs
index 9ffb8c2ee..9967376d1 100644
--- a/core/server/src/boot/listeners.rs
+++ b/core/server/src/boot/listeners.rs
@@ -403,11 +403,10 @@ pub(in crate::boot) fn make_replica_delegation_fns(
(accepted, dialed)
}
-/// Shard-0 client accept callbacks. TCP and WS clients are delegated via
-/// the coordinator (round-robin to peer shards); QUIC and TCP-TLS install
-/// locally on shard 0 because their per-connection state is not portable
-/// across shards (`compio_quic` endpoint binds one UDP socket; rustls TLS
-/// state ties to the post-handshake reactor).
+/// Shard-0 client accept callbacks. TCP, WS, TCP-TLS and WSS delegate
+/// raw sockets through the coordinator before any handshake. The
+/// destination shard supplies its handler and owns handshakes and I/O.
+/// The QUIC callback installs locally through shard 0's UDP endpoint.
// ws/wss bindings intentionally mirror the transport names (same convention as
// `replica_io::start_on_shard_zero`).
#[allow(clippy::similar_names)]
@@ -417,11 +416,7 @@ pub(in crate::boot) fn make_shard_zero_client_accept_fns(
on_request: RequestHandler,
) -> LocalClientAcceptFns {
let quic_bus = Rc::clone(bus);
- let tcp_tls_bus = Rc::clone(bus);
- let wss_bus = Rc::clone(bus);
- let quic_request = on_request.clone();
- let wss_request = on_request.clone();
- let tcp_tls_request = on_request;
+ let quic_request = on_request;
let tcp_coord = Rc::clone(&coord);
let tcp = Rc::new(move |stream| match tcp_coord.delegate_client(stream) {
@@ -435,9 +430,9 @@ pub(in crate::boot) fn make_shard_zero_client_accept_fns(
Err(error) => warn!(error = ?error, "delegate_ws_client failed;
dropping WS client"),
});
- // QUIC and TCP-TLS terminate locally on shard 0 but mint their client
+ // QUIC terminates locally on shard 0 but mints its client
// ids through the coordinator's `client_seq`, the same counter the
- // delegated TCP/WS path uses. A separate counter here would let a
+ // delegated TCP/WS/TCP-TLS/WSS paths use. A separate counter here would
let a
// shard-0-local id collide with a delegated id that round-robined to
// shard 0 (both encode target shard 0) in shard 0's connection
// registry.
@@ -449,30 +444,20 @@ pub(in crate::boot) fn make_shard_zero_client_accept_fns(
let tcp_tls_coord = Rc::clone(&coord);
let tcp_tls = Rc::new(move |stream, tls_config| {
- let Some(meta) =
- client_meta_from_stream(&stream, &tcp_tls_coord,
ClientTransportKind::TcpTls)
- else {
- return;
- };
- installer::install_client_tcp_tls(
- &tcp_tls_bus,
- meta,
- stream,
- tls_config,
- tcp_tls_request.clone(),
- );
+ match tcp_tls_coord.delegate_tcp_tls_client(stream, tls_config) {
+ Ok(client_id) => info!(client_id, "TCP-TLS client delegated"),
+ Err(error) => {
+ warn!(error = ?error, "delegate_tcp_tls_client failed;
dropping TCP-TLS client");
+ }
+ }
});
- // WSS terminates locally on shard 0 like TCP-TLS (rustls state is not
- // serialisable across the delegate path), minting ids through the same
- // coordinator counter.
let wss_coord = coord;
let wss = Rc::new(move |stream, tls_config| {
- let Some(meta) = client_meta_from_stream(&stream, &wss_coord,
ClientTransportKind::Wss)
- else {
- return;
- };
- installer::install_client_wss(&wss_bus, meta, stream, tls_config,
wss_request.clone());
+ match wss_coord.delegate_wss_client(stream, tls_config) {
+ Ok(client_id) => info!(client_id, "WSS client delegated"),
+ Err(error) => warn!(error = ?error, "delegate_wss_client failed;
dropping WSS client"),
+ }
});
LocalClientAcceptFns {
@@ -484,21 +469,6 @@ pub(in crate::boot) fn make_shard_zero_client_accept_fns(
}
}
-fn client_meta_from_stream(
- stream: &compio::net::TcpStream,
- coord: &shard::coordinator::ShardZeroCoordinator,
- transport: ClientTransportKind,
-) -> Option<ClientConnMeta> {
- let peer_addr = match stream.peer_addr() {
- Ok(peer_addr) => peer_addr,
- Err(error) => {
- warn!(error = %error, "dropping accepted client with unknown peer
address");
- return None;
- }
- };
- Some(mint_client_meta(coord, peer_addr, transport))
-}
-
fn mint_client_meta(
coord: &shard::coordinator::ShardZeroCoordinator,
peer_addr: SocketAddr,
diff --git a/core/server/src/boot/mod.rs b/core/server/src/boot/mod.rs
index e7a218695..36850f704 100644
--- a/core/server/src/boot/mod.rs
+++ b/core/server/src/boot/mod.rs
@@ -910,10 +910,9 @@ async fn shard_main(
}
// Listeners (replica + every client transport) bind on shard 0 only.
- // Shard 0's coordinator round-robins inbound TCP/WS connections to
- // peer shards via fd-transfer. QUIC and TCP-TLS clients terminate
- // locally on shard 0 (their per-connection state is non-portable -
- // see `LifecycleFrame::ClientWsConnectionSetup` rustdoc).
+ // The coordinator delegates TCP/WS/TCP-TLS/WSS before any handshake;
+ // the destination shard owns all connection state and I/O.
+ // QUIC terminates locally on shard 0 through its shared UDP endpoint.
if shard_id == 0 {
let coord = shard
.coordinator()
diff --git a/core/server/src/dispatch/mod.rs b/core/server/src/dispatch/mod.rs
index 8692bb630..cc2ac9c70 100644
--- a/core/server/src/dispatch/mod.rs
+++ b/core/server/src/dispatch/mod.rs
@@ -130,12 +130,11 @@ where
})
}
-/// Build the shard's one client-request handler: per-client FIFO queues
-/// drained one task per client, and the bus connection-lost hook that
-/// logs a dropped connection out. Every transport on the shard must
-/// share the instance (shard 0 hands it to its local QUIC, TCP-TLS and
-/// WSS listeners as well), or a client's ordering guarantee and the
-/// disconnect hook split by transport.
+/// Build the shard's client-request handler with per-client FIFO queues
+/// and a connection-lost hook. All connections installed on this shard
+/// share it to preserve ordering and disconnect cleanup across transports.
+/// The destination shard supplies it for delegated TCP/WS/TCP-TLS/WSS
+/// connections; shard 0 also supplies it for local QUIC connections.
pub fn make_deferred_client_request_handler<B, MJ, S, SB>(
bus: &B,
shard_handle: &ShellShardHandle<B, MJ, S, SB>,
diff --git a/core/server/src/dispatch/test_support.rs
b/core/server/src/dispatch/test_support.rs
index f00b64bca..c5bb3867f 100644
--- a/core/server/src/dispatch/test_support.rs
+++ b/core/server/src/dispatch/test_support.rs
@@ -30,7 +30,7 @@ use message_bus::installer::conn_info::ClientConnMeta;
use message_bus::replica::listener::MessageHandler;
use message_bus::{
BusMessage, ClientConnectionLostFn, ClientForwardFn, ConnectionLostFn,
JoinHandle, MessageBus,
- ReplicaForwardFn, ReplicaHandshakeDoneFn, SendError,
+ ReplicaForwardFn, ReplicaHandshakeDoneFn, SendError, SharedTlsServerConfig,
};
use metadata::impls::metadata::IggySnapshot;
use metadata::stm::stream::Streams;
@@ -144,6 +144,22 @@ impl ConnectionInstaller for SpyBus {
_on_request: RequestHandler,
) {
}
+ fn install_client_tcp_tls_fd(
+ &self,
+ _fd: DupedFd,
+ _meta: ClientConnMeta,
+ _config: SharedTlsServerConfig,
+ _on_request: RequestHandler,
+ ) {
+ }
+ fn install_client_wss_fd(
+ &self,
+ _fd: DupedFd,
+ _meta: ClientConnMeta,
+ _config: SharedTlsServerConfig,
+ _on_request: RequestHandler,
+ ) {
+ }
fn client_meta(&self, _client_id: u128) -> Option<Rc<ClientConnMeta>> {
None
}
diff --git a/core/shard/src/coordinator.rs b/core/shard/src/coordinator.rs
index 72de53324..04825edaf 100644
--- a/core/shard/src/coordinator.rs
+++ b/core/shard/src/coordinator.rs
@@ -25,7 +25,7 @@
//! 2. duplicates the TCP fd,
//! 3. sends a connection-setup `LifecycleFrame`
//! (`ReplicaInboundSetup` / `ReplicaOutboundSetup` /
-//! `Client{,Ws}ConnectionSetup`) to the target shard's inbox,
+//! `Client{,Ws,TcpTls,Wss}ConnectionSetup`) to the target shard's inbox,
//! 4. drops its own `TcpStream` so only the target shard's wrapped fd
//! keeps the socket alive.
//!
@@ -53,7 +53,7 @@ use crate::metrics::{frame_drop_reason, frame_drop_variant};
use crate::{LifecycleFrame, ShardCtorError, ShardFrame, TaggedSender,
validate_sender_ordering};
use compio::net::TcpStream;
use message_bus::installer::conn_info::{ClientConnMeta, ClientTransportKind};
-use message_bus::{SendError, fd_transfer};
+use message_bus::{SendError, SharedTlsServerConfig, fd_transfer};
use std::cell::Cell;
use std::rc::Rc;
use tracing::warn;
@@ -222,10 +222,10 @@ impl ShardZeroCoordinator {
}
/// Mint a client id for a connection that terminates locally on shard 0
- /// (QUIC, TCP-TLS) instead of being round-robin delegated.
+ /// (QUIC or HTTP) instead of being round-robin delegated.
///
/// Shares the coordinator's `client_seq` counter with the delegated
- /// TCP/WS path. Both a shard-0-local connection and a delegated
+ /// TCP/WS/TCP-TLS/WSS paths. Both a shard-0-local connection and a
delegated
/// connection that round-robined to shard 0 encode `target_shard = 0`
/// in the top 16 bits; drawing from separate counters would let them
/// mint the same id and collide in shard 0's single connection
@@ -323,23 +323,9 @@ impl ShardZeroCoordinator {
/// when the target shard's inbox refuses the setup frame (full or
/// disconnected).
pub fn delegate_client(&self, stream: TcpStream) -> Result<u128,
SendError> {
- let target = self.next_client_target();
- let client_id = self.mint_client_id(target);
- let peer_addr = stream.peer_addr().map_err(SendError::DupFailed)?;
-
- let fd = fd_transfer::dup_fd(&stream).map_err(SendError::DupFailed)?;
- let meta = ClientConnMeta::new(client_id, peer_addr,
ClientTransportKind::Tcp);
- let setup = LifecycleFrame::ClientConnectionSetup { fd, meta };
- if let Err(e) = self.senders[target as
usize].try_send(ShardFrame::lifecycle(setup)) {
- // The returned frame owns the `DupedFd` and closes it on drop.
- self.metrics
- .record_frame_drop(frame_drop_variant::FD_TRANSFER,
classify_try_send_err(&e));
- warn!(client_id, target, "delegate_client try_send failed: {e:?}");
- return Err(SendError::RoutingFailed(target));
- }
-
- drop(stream);
- Ok(client_id)
+ self.ship_client_fd(stream, ClientTransportKind::Tcp, |fd, meta| {
+ LifecycleFrame::ClientConnectionSetup { fd, meta }
+ })
}
/// Ship a WebSocket client's pre-upgrade TCP connection to the next
@@ -358,19 +344,71 @@ impl ShardZeroCoordinator {
/// fails or `dup(2)` fails. Returns [`SendError::RoutingFailed`]
/// when the target shard's inbox refuses the setup frame.
pub fn delegate_ws_client(&self, stream: TcpStream) -> Result<u128,
SendError> {
+ self.ship_client_fd(stream, ClientTransportKind::Ws, |fd, meta| {
+ LifecycleFrame::ClientWsConnectionSetup { fd, meta }
+ })
+ }
+
+ /// Delegate a raw TCP socket and its listener's TLS configuration.
+ /// The destination shard owns the TLS handshake and encrypted I/O.
+ ///
+ /// # Errors
+ ///
+ /// Returns the same lookup, duplication and routing errors as
+ /// [`Self::delegate_client`]. Both socket handles close on failure.
+ pub fn delegate_tcp_tls_client(
+ &self,
+ stream: TcpStream,
+ config: SharedTlsServerConfig,
+ ) -> Result<u128, SendError> {
+ self.ship_client_fd(stream, ClientTransportKind::TcpTls, |fd, meta| {
+ LifecycleFrame::ClientTcpTlsConnectionSetup { fd, meta, config }
+ })
+ }
+
+ /// Delegate WSS before TLS or WebSocket state is created. Both
+ /// handshakes and all subsequent I/O run on the destination shard.
+ ///
+ /// # Errors
+ ///
+ /// Returns the same lookup, duplication and routing errors as
+ /// [`Self::delegate_client`]. Both socket handles close on failure.
+ pub fn delegate_wss_client(
+ &self,
+ stream: TcpStream,
+ config: SharedTlsServerConfig,
+ ) -> Result<u128, SendError> {
+ self.ship_client_fd(stream, ClientTransportKind::Wss, |fd, meta| {
+ LifecycleFrame::ClientWssConnectionSetup { fd, meta, config }
+ })
+ }
+
+ #[must_use]
+ pub const fn total_shards(&self) -> u16 {
+ self.total_shards
+ }
+
+ fn ship_client_fd(
+ &self,
+ stream: TcpStream,
+ transport: ClientTransportKind,
+ build_frame: impl FnOnce(fd_transfer::DupedFd, ClientConnMeta) ->
LifecycleFrame,
+ ) -> Result<u128, SendError> {
let target = self.next_client_target();
let client_id = self.mint_client_id(target);
let peer_addr = stream.peer_addr().map_err(SendError::DupFailed)?;
let fd = fd_transfer::dup_fd(&stream).map_err(SendError::DupFailed)?;
- let meta = ClientConnMeta::new(client_id, peer_addr,
ClientTransportKind::Ws);
- let setup = LifecycleFrame::ClientWsConnectionSetup { fd, meta };
+ let meta = ClientConnMeta::new(client_id, peer_addr, transport);
+ let setup = build_frame(fd, meta);
if let Err(e) = self.senders[target as
usize].try_send(ShardFrame::lifecycle(setup)) {
self.metrics
.record_frame_drop(frame_drop_variant::FD_TRANSFER,
classify_try_send_err(&e));
warn!(
client_id,
- target, "delegate_ws_client try_send failed: {e:?}"
+ target,
+ ?transport,
+ "client delegation try_send failed: {e:?}"
);
return Err(SendError::RoutingFailed(target));
}
@@ -378,11 +416,6 @@ impl ShardZeroCoordinator {
drop(stream);
Ok(client_id)
}
-
- #[must_use]
- pub const fn total_shards(&self) -> u16 {
- self.total_shards
- }
}
/// Map `crossfire::TrySendError` to the `frame_drop_reason` label set so
@@ -412,7 +445,21 @@ fn rr_pick(counter: &Cell<u16>, total_shards: u16,
skip_zero: bool) -> u16 {
#[cfg(test)]
mod tests {
use super::*;
+ use compio::io::AsyncRead;
use compio::net::{TcpListener, TcpStream};
+ use message_bus::client_listener::tcp_tls;
+ use message_bus::transports::tls::self_signed_for_loopback;
+ use std::collections::HashSet;
+ use std::sync::Arc;
+ use std::time::Duration;
+
+ const CLIENT_TRANSPORTS: [ClientTransportKind; 4] = [
+ ClientTransportKind::Tcp,
+ ClientTransportKind::Ws,
+ ClientTransportKind::TcpTls,
+ ClientTransportKind::Wss,
+ ];
+ const SOCKET_CLOSE_TIMEOUT: Duration = Duration::from_secs(2);
fn build_senders(total: u16) -> Rc<Vec<TaggedSender>> {
let mut senders = Vec::with_capacity(total as usize);
@@ -723,7 +770,7 @@ mod tests {
)
.expect("coord ctor ok");
- // Interleave shard-0-local mints (QUIC / TCP-TLS path) with
+ // Interleave shard-0-local mints (QUIC / HTTP path) with
// delegated mints that round-robin to shard 0. Both encode target
// shard 0 in the top 16 bits; the shared counter keeps every id
// distinct so they cannot overwrite each other in shard 0's
@@ -898,4 +945,205 @@ mod tests {
_ => panic!("expected ClientWsConnectionSetup variant"),
}
}
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn mixed_client_transports_share_placement_and_recovered_sequence() {
+ const RECOVERED_SEQUENCE: u128 = 42;
+ const ACCEPTS: u16 = 16;
+ let configs = [test_tls_config(), test_tls_config()];
+ for total_shards in [1, 3, 4] {
+ for skip_zero in [false, true] {
+ let (senders, receivers) = build_senders_with_rx(total_shards);
+ let coord = ShardZeroCoordinator::new(
+ senders,
+ total_shards,
+ CoordinatorConfig {
+ skip_shard_zero_for_clients: skip_zero,
+ ..CoordinatorConfig::default()
+ },
+ crate::metrics::ShardMetrics::for_shard(),
+ )
+ .unwrap();
+ coord.seed_client_sequence(0,
std::iter::once(RECOVERED_SEQUENCE));
+ let first_shard = u16::from(skip_zero && total_shards > 1);
+ let mut sequences = HashSet::new();
+
+ for index in 0..ACCEPTS {
+ let transport = CLIENT_TRANSPORTS[usize::from(index) %
CLIENT_TRANSPORTS.len()];
+ let config =
+ &configs[(usize::from(index) /
CLIENT_TRANSPORTS.len()) % configs.len()];
+ let (stream, mut peer) = tcp_pair().await;
+ let peer_addr = peer.local_addr().unwrap();
+ let client_id =
+ delegate_test_client(&coord, stream, transport,
config).unwrap();
+ let target = first_shard + index % (total_shards -
first_shard);
+ assert_eq!(message_bus::client_id_owning_shard(client_id),
target);
+ let sequence = client_id & CLIENT_SEQUENCE_MASK;
+ assert_eq!(sequence, RECOVERED_SEQUENCE +
u128::from(index) * 2 + 1);
+ assert!(sequences.insert(sequence));
+
assert!(sequences.insert(coord.mint_shard_zero_client_id()));
+
+ let frame =
receivers[usize::from(target)].try_recv().unwrap();
+ let (fd, meta) = client_setup(frame, transport, config);
+ assert_eq!(meta.client_id, client_id);
+ assert_eq!(meta.transport, transport);
+ assert_eq!(meta.peer_addr, peer_addr);
+ assert!(
+ receivers
+ .iter()
+ .all(|receiver| receiver.try_recv().is_err())
+ );
+
+ drop(fd);
+ assert_socket_closed(&mut peer).await;
+ }
+ }
+ }
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn rejected_client_handoffs_close_both_socket_handles() {
+ let config = test_tls_config();
+ for transport in CLIENT_TRANSPORTS {
+ for disconnected in [false, true] {
+ let (senders, receivers) = build_senders_with_rx(2);
+ if !disconnected {
+ while senders[1]
+ .try_send(ShardFrame::lifecycle(
+ LifecycleFrame::ReplicaInboundHandshakeDone {
slot: 0 },
+ ))
+ .is_ok()
+ {}
+ }
+ let _receivers = if disconnected {
+ drop(receivers);
+ None
+ } else {
+ Some(receivers)
+ };
+ let coord = ShardZeroCoordinator::new(
+ senders,
+ 2,
+ CoordinatorConfig {
+ skip_shard_zero_for_clients: true,
+ ..CoordinatorConfig::default()
+ },
+ crate::metrics::ShardMetrics::for_shard(),
+ )
+ .unwrap();
+ let (stream, mut peer) = tcp_pair().await;
+ let result = delegate_test_client(&coord, stream, transport,
&config);
+ assert!(
+ matches!(result, Err(SendError::RoutingFailed(1))),
+ "{transport:?}, disconnected={disconnected}: {result:?}"
+ );
+ assert_eq!(Arc::strong_count(&config), 1);
+ assert_socket_closed(&mut peer).await;
+ }
+ }
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn
dropping_queued_tls_setups_releases_sockets_and_listener_configs() {
+ let config = test_tls_config();
+ for transport in [ClientTransportKind::TcpTls,
ClientTransportKind::Wss] {
+ let (senders, receivers) = build_senders_with_rx(1);
+ let coord = ShardZeroCoordinator::new(
+ senders,
+ 1,
+ CoordinatorConfig::default(),
+ crate::metrics::ShardMetrics::for_shard(),
+ )
+ .unwrap();
+ let (stream, mut peer) = tcp_pair().await;
+ delegate_test_client(&coord, stream, transport, &config).unwrap();
+ assert_eq!(Arc::strong_count(&config), 2);
+ drop(coord);
+ drop(receivers);
+ assert_eq!(Arc::strong_count(&config), 1);
+ assert_socket_closed(&mut peer).await;
+ }
+ }
+
+ fn test_tls_config() -> SharedTlsServerConfig {
+ tcp_tls::bind("127.0.0.1:0".parse().unwrap(),
self_signed_for_loopback())
+ .unwrap()
+ .1
+ }
+
+ #[allow(clippy::future_not_send)]
+ async fn tcp_pair() -> (TcpStream, TcpStream) {
+ let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let (connected, accepted) = futures::join!(
+ TcpStream::connect(listener.local_addr().unwrap()),
+ listener.accept()
+ );
+ (accepted.unwrap().0, connected.unwrap())
+ }
+
+ fn delegate_test_client(
+ coord: &ShardZeroCoordinator,
+ stream: TcpStream,
+ transport: ClientTransportKind,
+ config: &SharedTlsServerConfig,
+ ) -> Result<u128, SendError> {
+ match transport {
+ ClientTransportKind::Tcp => coord.delegate_client(stream),
+ ClientTransportKind::Ws => coord.delegate_ws_client(stream),
+ ClientTransportKind::TcpTls => {
+ coord.delegate_tcp_tls_client(stream, Arc::clone(config))
+ }
+ ClientTransportKind::Wss => coord.delegate_wss_client(stream,
Arc::clone(config)),
+ _ => panic!("transport has no TCP delegation path: {transport:?}"),
+ }
+ }
+
+ fn client_setup(
+ frame: ShardFrame,
+ transport: ClientTransportKind,
+ expected_config: &SharedTlsServerConfig,
+ ) -> (fd_transfer::DupedFd, ClientConnMeta) {
+ match (transport, frame) {
+ (
+ ClientTransportKind::Tcp,
+ ShardFrame::Lifecycle(LifecycleFrame::ClientConnectionSetup {
fd, meta }),
+ )
+ | (
+ ClientTransportKind::Ws,
+ ShardFrame::Lifecycle(LifecycleFrame::ClientWsConnectionSetup
{ fd, meta }),
+ ) => (fd, meta),
+ (
+ ClientTransportKind::TcpTls,
+
ShardFrame::Lifecycle(LifecycleFrame::ClientTcpTlsConnectionSetup {
+ fd,
+ meta,
+ config,
+ }),
+ )
+ | (
+ ClientTransportKind::Wss,
+ ShardFrame::Lifecycle(LifecycleFrame::ClientWssConnectionSetup
{
+ fd,
+ meta,
+ config,
+ }),
+ ) => {
+ assert!(Arc::ptr_eq(&config, expected_config));
+ (fd, meta)
+ }
+ _ => panic!("wrong setup variant for {transport:?}"),
+ }
+ }
+
+ #[allow(clippy::future_not_send)]
+ async fn assert_socket_closed(peer: &mut TcpStream) {
+ let result = compio::time::timeout(SOCKET_CLOSE_TIMEOUT,
peer.read(vec![0]))
+ .await
+ .expect("delegation must release both socket handles")
+ .0;
+ assert_eq!(result.unwrap(), 0, "peer must observe EOF");
+ }
}
diff --git a/core/shard/src/lib.rs b/core/shard/src/lib.rs
index 90121ace8..f85c29ecd 100644
--- a/core/shard/src/lib.rs
+++ b/core/shard/src/lib.rs
@@ -58,7 +58,7 @@ use message_bus::client_listener::RequestHandler;
use message_bus::fd_transfer::DupedFd;
use message_bus::installer::conn_info::{ClientConnMeta, ClientTransportKind};
use message_bus::replica::listener::MessageHandler;
-use message_bus::{BusMessage, MessageBus};
+use message_bus::{BusMessage, MessageBus, SharedTlsServerConfig};
use metadata::IggyMetadata;
use metadata::impls::metadata::StreamsFrontend;
use metadata::stm::StateMachine;
@@ -715,7 +715,7 @@ pub enum LifecycleFrame {
/// Shard 0 distributes an inbound SDK WebSocket client's pre-upgrade
/// TCP connection fd to the owning shard. The HTTP-Upgrade handshake
/// has NOT run yet at this point: the fd is plain TCP, the dup is
- /// safe (cross-shard fd-delegation only happens for plain TCP), and
+ /// safe before any transport state is created, and
/// `compio_ws::WebSocketStream<TcpStream>`'s `!Send` constraint
/// (compio `Rc<...>` driver state, post-upgrade) does not apply.
/// The receiving shard wraps the fd, runs `compio_ws::accept_async`,
@@ -734,6 +734,21 @@ pub enum LifecycleFrame {
fd: DupedFd,
meta: ClientConnMeta,
},
+ /// Delegate TCP-TLS before reading TLS bytes. The destination wraps the
+ /// fd on its runtime and owns the handshake and connection tasks.
+ ClientTcpTlsConnectionSetup {
+ fd: DupedFd,
+ meta: ClientConnMeta,
+ config: SharedTlsServerConfig,
+ },
+ /// Delegate WSS before either handshake. The listener's configuration
+ /// travels with the socket; all TLS and WebSocket state stays local to
+ /// the destination runtime.
+ ClientWssConnectionSetup {
+ fd: DupedFd,
+ meta: ClientConnMeta,
+ config: SharedTlsServerConfig,
+ },
/// A non-owning shard forwards a replica send to the owning shard's
/// local bus; the owning shard then takes the fast path.
ForwardReplicaSend {
@@ -894,6 +909,10 @@ pub enum ShardFrame {
// rather than as a failing test.
const _: () = assert!(std::mem::size_of::<ShardFrame>() ==
std::mem::size_of::<LifecycleFrame>());
+// Every inbox slot pays for the largest variant, including consensus traffic.
+const MAX_SHARD_FRAME_SIZE: usize = 160;
+const _: () = assert!(std::mem::size_of::<ShardFrame>() <=
MAX_SHARD_FRAME_SIZE);
+
impl ShardFrame {
/// Create a consensus frame addressed to `target_shard`. The sender
/// is the routing authority; `accept_frame_for_self` compares this
diff --git a/core/shard/src/router.rs b/core/shard/src/router.rs
index dccf089b8..9dec7e986 100644
--- a/core/shard/src/router.rs
+++ b/core/shard/src/router.rs
@@ -747,6 +747,30 @@ where
self.bus
.install_client_ws_fd(fd, meta,
self.on_client_request.clone());
}
+ LifecycleFrame::ClientTcpTlsConnectionSetup { fd, meta, config }
=> {
+ tracing::info!(
+ shard = self.id,
+ client_id = meta.client_id,
+ raw_fd = fd.as_raw_fd(),
+ "installing delegated TCP-TLS client fd (pre-handshake)"
+ );
+ self.bus.install_client_tcp_tls_fd(
+ fd,
+ meta,
+ config,
+ self.on_client_request.clone(),
+ );
+ }
+ LifecycleFrame::ClientWssConnectionSetup { fd, meta, config } => {
+ tracing::info!(
+ shard = self.id,
+ client_id = meta.client_id,
+ raw_fd = fd.as_raw_fd(),
+ "installing delegated WSS client fd (pre-handshake)"
+ );
+ self.bus
+ .install_client_wss_fd(fd, meta, config,
self.on_client_request.clone());
+ }
LifecycleFrame::ForwardReplicaSend { replica_id, msg } => {
if let Err(e) = self.bus.send_to_replica(replica_id,
msg).await {
tracing::debug!(
diff --git a/core/simulator/src/bus.rs b/core/simulator/src/bus.rs
index 7002ea4a1..0434107ed 100644
--- a/core/simulator/src/bus.rs
+++ b/core/simulator/src/bus.rs
@@ -27,7 +27,7 @@ use message_bus::installer::conn_info::ClientConnMeta;
use message_bus::replica::listener::MessageHandler;
use message_bus::{
BusMessage, ClientConnectionLostFn, ConnectionInstaller, MessageBus,
ReplicaHandshakeDoneFn,
- SendError,
+ SendError, SharedTlsServerConfig,
};
use server_common::{
MESSAGE_ALIGN, Message,
@@ -368,6 +368,26 @@ impl ConnectionInstaller for SharedSimOutbox {
panic!("simulator has no fd transfer: ws client install is
unreachable");
}
+ fn install_client_tcp_tls_fd(
+ &self,
+ _fd: DupedFd,
+ _meta: ClientConnMeta,
+ _config: SharedTlsServerConfig,
+ _on_request: RequestHandler,
+ ) {
+ panic!("simulator has no fd transfer: tcp tls client install is
unreachable");
+ }
+
+ fn install_client_wss_fd(
+ &self,
+ _fd: DupedFd,
+ _meta: ClientConnMeta,
+ _config: SharedTlsServerConfig,
+ _on_request: RequestHandler,
+ ) {
+ panic!("simulator has no fd transfer: wss client install is
unreachable");
+ }
+
// Real, unlike the fd installs above: dispatch reads client_meta to seed
// the session manager, and registers the disconnect callback the harness
// fires. Both delegate to the in-memory registry on the inner outbox.
