You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iggy.apache.org by hg...@apache.org on 2026/04/17 14:20:04 UTC
(iggy) branch bus updated (2bf078291 -> d3d93fa33)
This is an automated email from the ASF dual-hosted git repository.
hgruszecki pushed a change to branch bus
in repository https://gitbox.apache.org/repos/asf/iggy.git
omit 2bf078291 feat(message_bus): async fire-and-forget transport for VSR consensus
new d3d93fa33 feat(message_bus): async fire-and-forget transport for VSR consensus
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (2bf078291)
\
N -- N -- N refs/heads/bus (d3d93fa33)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
core/binary_protocol/src/consensus/iobuf.rs | 5 +-
core/message_bus/src/config.rs | 88 +++++
core/message_bus/src/connector.rs | 50 ++-
core/message_bus/src/error.rs | 4 +-
core/message_bus/src/fd_transfer.rs | 59 +++-
core/message_bus/src/framing.rs | 23 +-
core/message_bus/src/installer.rs | 162 +++++++--
core/message_bus/src/lib.rs | 138 +++++---
.../src/lifecycle/connection_registry.rs | 94 +++--
core/message_bus/src/lifecycle/mod.rs | 4 +-
core/message_bus/src/replica_io.rs | 10 +-
core/message_bus/src/replica_listener.rs | 23 +-
core/message_bus/src/socket_opts.rs | 32 +-
core/message_bus/src/writer_task.rs | 19 +-
core/message_bus/tests/backpressure.rs | 11 +-
core/message_bus/tests/connection_lost_notify.rs | 122 +++++++
core/message_bus/tests/directional_connection.rs | 22 +-
core/message_bus/tests/duplicate_client_id.rs | 84 ++++-
core/message_bus/tests/graceful_shutdown.rs | 58 +++
core/message_bus/tests/head_of_line.rs | 22 +-
core/message_bus/tests/reconnect.rs | 11 +-
core/message_bus/tests/reconnect_skip_connected.rs | 110 ++++++
core/message_bus/tests/replica_roundtrip.rs | 11 +-
core/message_bus/tests/shard_zero_gating.rs | 11 +-
core/message_bus/tests/tcp_client_roundtrip.rs | 2 +-
core/message_bus/tests/vectored_batch.rs | 11 +-
core/partitions/src/journal.rs | 47 +++
core/shard/src/config.rs | 60 ++++
core/shard/src/coordinator.rs | 387 +++++++++++++++++----
core/shard/src/lib.rs | 163 ++++++++-
core/shard/src/router.rs | 12 +-
31 files changed, 1600 insertions(+), 255 deletions(-)
create mode 100644 core/message_bus/src/config.rs
create mode 100644 core/message_bus/tests/connection_lost_notify.rs
create mode 100644 core/message_bus/tests/reconnect_skip_connected.rs
create mode 100644 core/shard/src/config.rs
(iggy) 01/01: feat(message_bus): async fire-and-forget transport for VSR consensus
Posted by hg...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
hgruszecki pushed a commit to branch bus
in repository https://gitbox.apache.org/repos/asf/iggy.git
commit d3d93fa3302f23fa9e2657c1a48f3c59ec49281e
Author: Hubert Gruszecki <h....@gmail.com>
AuthorDate: Fri Apr 17 15:24:53 2026 +0200
feat(message_bus): async fire-and-forget transport for VSR consensus
Adds the message_bus crate: a compio + io_uring TCP transport that
backs the VSR consensus plane. Replicas and SDK clients both flow
through per-peer bounded mpsc queues drained by dedicated writer
tasks that coalesce up to MAX_BATCH messages into a single writev
syscall. A slow peer cannot stall sends to other peers; VSR's WAL /
timeout machinery covers drops.
Connection lifecycle, fd transfer, and shard-0 routing:
- Shard 0 is the sole binder of the replica / client listeners and
the sole outbound dialer for higher-id peers. On each accept or
successful dial the ShardZeroCoordinator dups the fd via
fcntl(F_DUPFD_CLOEXEC), ships it to the target shard via an
inter-shard ShardFrame channel, and drops its own stream so only
the target's wrapped fd keeps the socket alive.
- TaggedSender<R> newtype lifts the senders[i].shard_id() == i
invariant from a doc comment to a ctor assertion
(assert_sender_ordering). Permuted sender Vecs now fail at
bootstrap instead of silently misrouting every setup / mapping
/ forward frame.
- Authoritative replica_id -> owning_shard snapshot lives on the
coordinator as [Option<u16>; 256]. A periodic refresh task
re-broadcasts the snapshot so shards whose inbox was full when
the original ReplicaMappingUpdate went out recover on the next
tick. ConnectionLost handler invokes broadcast_mapping_clear
cluster-wide.
- install_{replica,client}_stream reserves the registry key before
spawning reader/writer tasks and aborts both halves (in-band
install_aborted Cell) when the insert loses a duplicate-key
race. Previously the losing reader stayed alive and could
clobber the winner's mapping via notify_connection_lost or
misroute on_request via the wrong client_id.
- notify_connection_lost fires exactly once per disconnect via a
shared one-shot Cell<bool> guard across the writer and reader
halves.
- periodic_reconnect skips peers already present in
bus.replicas() or mapped via owning_shard. The dialer's contains
check at install time fired AFTER the waste; reconnect storms
against already-connected peers are eliminated.
- Partition retransmit sends Frozen directly from the journal
instead of parse-reparse + per-target 4 KiB memcpy. The
.expect("partition journal entry must contain valid prepare")
that would panic the shard on a corrupted journal entry is gone.
- ConnectionLost handler on shard 0 calls
coordinator.forget_mapping + broadcast_mapping_clear before
touching local state, so other shards stop routing via a dead
replica immediately.
Transport knobs:
- MessageBusConfig (max_batch, max_message_size,
peer_queue_capacity, reconnect_period, keepalive_idle /
_interval / _retries, close_peer_timeout) replaces the set of
hardcoded module-level consts. Threads through
IggyMessageBus::new / ::with_config, writer_task::run,
framing::read_message, socket_opts::
apply_keepalive_for_connection, installer, and connector.
- CoordinatorConfig (refresh_period plus two shard-0 skip flags)
replaces the hardcoded refresh cadence and the prior
shard-0-always RR. Default: skip shard 0 for replicas (long-lived
steady flows offload to peer shards), keep shard 0 in the client
RR (short-lived flows benefit from shard-0 parallelism). Both
flags are configurable. rr_pick handles the skip logic and
degrades safely when total_shards == 1.
- Both config types live in-crate under core/message_bus/src/config.rs
and core/shard/src/config.rs with TODO markers pointing to their
eventual home under core/configs once downstream server bootstrap
that constructs the transport from ServerConfig lands.
MAX_MESSAGE_SIZE and DEFAULT_RECONNECT_PERIOD retained as pub
consts for test / bench ergonomics; other consts deleted.
Lifecycle and send-path correctness:
- Send-path borrow: ConnectionRegistry::borrow_sender and
ReplicaRegistry::borrow_sender return Option<Ref<'_, BusSender>>
via Ref::filter_map instead of a cloned Sender. try_send takes
&self, so the prior clone was an atomic RMW per send on the
inner Arc<State>. Zero atomic RMW on the fast path now.
- Graceful shutdown drains clients + replicas before background
tasks (accept loops, reconnect periodic, coordinator refresh).
A misbehaving background task can no longer consume the full
timeout budget and force writer tasks to tear down mid
write_vectored_all.
- drain_entries splits the per-entry deadline. Writer gets the
full shared deadline; reader gets at least READER_DRAIN_FLOOR
(250 ms) or half of what remains after the writer returns. Slow
writers no longer force-cancel the reader.
- IggyMessageBus::set_replica_forward_fn / set_client_forward_fn
take &self + RefCell<Option<ShardForwardFn>>, matching the
sibling set_connection_lost_fn / set_shard_mapping pattern.
Production wraps the bus in Rc, so the prior &mut self
signature was unreachable post-wrap.
Security + robustness:
- dup_fd uses fcntl(F_DUPFD_CLOEXEC) so no future fork+exec child
can inherit replica or client sockets. close_fd logs non-EINTR
failures that previously went silent.
- SendError::DupFailed payload is i32 (errno), not String. The
only heap-allocating variant is on a resource-exhaustion path
and is now gone.
- replica_listener::handshake carries a TODO(IGGY-112) marker:
current check validates cluster_id and directional id bound
only; no shared secret / mTLS / version negotiation.
Acceptable on the assumed trusted network; hardening tracked
under IGGY-112.
Tests:
- Integration tests across tcp_client_roundtrip, replica_roundtrip,
directional_connection, reconnect, reconnect_skip_connected,
shard_zero_gating, head_of_line, vectored_batch, backpressure,
connection_lost_notify, duplicate_client_id, graceful_shutdown.
- Unit tests for framing, connection_registry drain budgets,
coordinator RR with the four skip-flag combinations,
fd_transfer FD_CLOEXEC assertion, and ShardZeroCoordinator
sender-ordering panic.
cargo build, cargo fmt --all, cargo clippy --workspace --all-targets,
and cargo nextest run --workspace (2284 passed, 96 skipped) all green.
---
Cargo.lock | 117 ++--
DEPENDENCIES.md | 32 +-
core/binary_protocol/src/consensus/iobuf.rs | 64 +-
core/binary_protocol/src/consensus/message.rs | 2 +-
core/binary_protocol/src/consensus/mod.rs | 4 +-
core/configs/src/server_config/cluster.rs | 18 +
core/configs/src/server_config/defaults.rs | 4 +
core/configs/src/server_config/validators.rs | 123 ++++
core/consensus/src/impls.rs | 18 +-
core/consensus/src/plane_helpers.rs | 138 ++---
core/message_bus/Cargo.toml | 9 +-
core/message_bus/src/cache/connection.rs | 437 ++++++++++---
core/message_bus/src/cache/mod.rs | 3 +-
core/message_bus/src/client_listener.rs | 97 +++
core/message_bus/src/config.rs | 88 +++
core/message_bus/src/connector.rs | 184 ++++++
core/message_bus/src/{cache/mod.rs => error.rs} | 40 +-
core/message_bus/src/fd_transfer.rs | 125 ++++
core/message_bus/src/framing.rs | 209 +++++++
core/message_bus/src/installer.rs | 396 ++++++++++++
core/message_bus/src/lib.rs | 636 ++++++++++++++++---
.../src/lifecycle/connection_registry.rs | 682 +++++++++++++++++++++
core/message_bus/src/{cache => lifecycle}/mod.rs | 23 +-
core/message_bus/src/lifecycle/shutdown.rs | 184 ++++++
core/message_bus/src/replica_io.rs | 154 +++++
core/message_bus/src/replica_listener.rs | 164 +++++
core/message_bus/src/socket_opts.rs | 59 ++
core/message_bus/src/{cache => transports}/mod.rs | 18 +-
core/message_bus/src/writer_task.rs | 100 +++
core/message_bus/tests/backpressure.rs | 114 ++++
core/message_bus/tests/common/mod.rs | 88 +++
core/message_bus/tests/connection_lost_notify.rs | 122 ++++
core/message_bus/tests/directional_connection.rs | 116 ++++
core/message_bus/tests/duplicate_client_id.rs | 157 +++++
core/message_bus/tests/graceful_shutdown.rs | 134 ++++
core/message_bus/tests/head_of_line.rs | 151 +++++
core/message_bus/tests/reconnect.rs | 80 +++
core/message_bus/tests/reconnect_skip_connected.rs | 110 ++++
core/message_bus/tests/replica_roundtrip.rs | 127 ++++
core/message_bus/tests/shard_zero_gating.rs | 152 +++++
core/message_bus/tests/tcp_client_roundtrip.rs | 110 ++++
core/message_bus/tests/vectored_batch.rs | 107 ++++
core/metadata/src/impls/metadata.rs | 35 +-
core/partitions/src/iggy_partition.rs | 26 +-
core/partitions/src/iggy_partitions.rs | 7 +-
core/partitions/src/journal.rs | 47 ++
core/server/config.toml | 5 +-
core/shard/Cargo.toml | 1 +
core/shard/src/config.rs | 60 ++
core/shard/src/coordinator.rs | 561 +++++++++++++++++
core/shard/src/lib.rs | 333 ++++++++--
core/shard/src/router.rs | 158 ++++-
core/simulator/src/bus.rs | 115 ++--
core/simulator/src/lib.rs | 3 +-
core/simulator/src/replica.rs | 1 +
55 files changed, 6503 insertions(+), 545 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index f1a4280d5..14e9ef928 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -87,7 +87,7 @@ dependencies = [
"mime",
"percent-encoding",
"pin-project-lite",
- "rand 0.9.4",
+ "rand 0.9.2",
"sha1",
"smallvec",
"tokio",
@@ -416,7 +416,7 @@ dependencies = [
"miniz_oxide",
"num-bigint",
"quad-rand",
- "rand 0.9.4",
+ "rand 0.9.2",
"regex-lite",
"serde",
"serde_bytes",
@@ -1436,7 +1436,7 @@ dependencies = [
"log",
"num",
"pin-project-lite",
- "rand 0.9.4",
+ "rand 0.9.2",
"rustls",
"rustls-native-certs",
"rustls-pki-types",
@@ -1569,7 +1569,7 @@ dependencies = [
"indexmap 2.14.0",
"js-sys",
"once_cell",
- "rand 0.9.4",
+ "rand 0.9.2",
"serde",
"serde_bytes",
"serde_json",
@@ -1751,9 +1751,9 @@ dependencies = [
[[package]]
name = "cc"
-version = "1.2.60"
+version = "1.2.59"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "43c5703da9466b66a946814e1adf53ea2c90f10063b86290cc9eb67ce3478a20"
+checksum = "b7a4d3ec6524d28a329fc53654bbadc9bdd7b0431f5d65f1a56ffb28a1ee5283"
dependencies = [
"find-msvc-tools",
"jobserver",
@@ -1796,7 +1796,7 @@ checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601"
dependencies = [
"cfg-if",
"cpufeatures 0.3.0",
- "rand_core 0.10.1",
+ "rand_core 0.10.0",
]
[[package]]
@@ -3694,7 +3694,7 @@ checksum = "4e7f34442dbe69c60fe8eaf58a8cafff81a1f278816d8ab4db255b3bef4ac3c4"
dependencies = [
"getrandom 0.3.4",
"libm",
- "rand 0.9.4",
+ "rand 0.9.2",
"siphasher",
]
@@ -3755,7 +3755,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bb330bbd4cb7a5b9f559427f06f98a4f853a137c8298f3bd3f8ca57663e21986"
dependencies = [
"portable-atomic",
- "rand 0.9.4",
+ "rand 0.9.2",
"web-time",
]
@@ -4200,7 +4200,7 @@ dependencies = [
"js-sys",
"libc",
"r-efi 6.0.0",
- "rand_core 0.10.1",
+ "rand_core 0.10.0",
"wasip2",
"wasip3",
"wasm-bindgen",
@@ -4678,7 +4678,7 @@ dependencies = [
"parking_lot",
"portable-atomic",
"quanta",
- "rand 0.9.4",
+ "rand 0.9.2",
"smallvec",
"spinning_top",
"web-time",
@@ -4904,7 +4904,7 @@ dependencies = [
"idna",
"ipnet",
"once_cell",
- "rand 0.9.4",
+ "rand 0.9.2",
"ring",
"thiserror 2.0.18",
"tinyvec",
@@ -4926,7 +4926,7 @@ dependencies = [
"moka",
"once_cell",
"parking_lot",
- "rand 0.9.4",
+ "rand 0.9.2",
"resolv-conf",
"smallvec",
"thiserror 2.0.18",
@@ -5128,9 +5128,9 @@ dependencies = [
[[package]]
name = "hyper-rustls"
-version = "0.27.8"
+version = "0.27.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c2b52f86d1d4bc0d6b4e6826d960b1b333217e07d36b882dca570a5e1c48895b"
+checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58"
dependencies = [
"http 1.4.0",
"hyper",
@@ -5138,6 +5138,7 @@ dependencies = [
"log",
"rustls",
"rustls-native-certs",
+ "rustls-pki-types",
"tokio",
"tokio-rustls",
"tower-service",
@@ -6406,9 +6407,9 @@ dependencies = [
[[package]]
name = "js-sys"
-version = "0.3.95"
+version = "0.3.94"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2964e92d1d9dc3364cae4d718d93f227e3abb088e747d92e0395bfdedf1c12ca"
+checksum = "2e04e2ef80ce82e13552136fabeef8a5ed1f985a96805761cbb9a2c34e7664d9"
dependencies = [
"cfg-if",
"futures-util",
@@ -7007,9 +7008,16 @@ dependencies = [
name = "message_bus"
version = "0.1.0"
dependencies = [
+ "async-channel",
+ "compio",
+ "futures",
"iggy_binary_protocol",
"iggy_common",
+ "libc",
"rand 0.10.1",
+ "socket2 0.6.3",
+ "thiserror 2.0.18",
+ "tracing",
]
[[package]]
@@ -7197,7 +7205,7 @@ dependencies = [
"mongodb-internal-macros",
"pbkdf2",
"percent-encoding",
- "rand 0.9.4",
+ "rand 0.9.2",
"rustc_version_runtime",
"rustls",
"rustversion",
@@ -7672,9 +7680,9 @@ dependencies = [
[[package]]
name = "openssl"
-version = "0.10.77"
+version = "0.10.76"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bfe4646e360ec77dff7dde40ed3d6c5fee52d156ef4a62f53973d38294dad87f"
+checksum = "951c002c75e16ea2c65b8c7e4d3d51d5530d8dfa7d060b4776828c88cfb18ecf"
dependencies = [
"bitflags 2.11.0",
"cfg-if",
@@ -7713,9 +7721,9 @@ dependencies = [
[[package]]
name = "openssl-sys"
-version = "0.9.113"
+version = "0.9.112"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ad2f2c0eba47118757e4c6d2bff2838f3e0523380021356e7875e858372ce644"
+checksum = "57d55af3b3e226502be1526dfdba67ab0e9c96fc293004e79576b2b9edb0dbdb"
dependencies = [
"cc",
"libc",
@@ -7813,7 +7821,7 @@ dependencies = [
"futures-util",
"opentelemetry",
"percent-encoding",
- "rand 0.9.4",
+ "rand 0.9.2",
"thiserror 2.0.18",
"tokio",
"tokio-stream",
@@ -8249,9 +8257,9 @@ dependencies = [
[[package]]
name = "pkg-config"
-version = "0.3.33"
+version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e"
+checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c"
[[package]]
name = "plain"
@@ -8732,7 +8740,7 @@ dependencies = [
"fastbloom",
"getrandom 0.3.4",
"lru-slab",
- "rand 0.9.4",
+ "rand 0.9.2",
"ring",
"rustc-hash",
"rustls",
@@ -8799,9 +8807,9 @@ dependencies = [
[[package]]
name = "rand"
-version = "0.9.4"
+version = "0.9.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea"
+checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1"
dependencies = [
"rand_chacha 0.9.0",
"rand_core 0.9.5",
@@ -8815,7 +8823,7 @@ checksum = "d2e8e8bcc7961af1fdac401278c6a831614941f6164ee3bf4ce61b7edb162207"
dependencies = [
"chacha20",
"getrandom 0.4.2",
- "rand_core 0.10.1",
+ "rand_core 0.10.0",
]
[[package]]
@@ -8858,9 +8866,9 @@ dependencies = [
[[package]]
name = "rand_core"
-version = "0.10.1"
+version = "0.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69"
+checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba"
[[package]]
name = "rand_xoshiro"
@@ -8868,7 +8876,7 @@ version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1f0b2cc7bfeef8f0320ca45f88b00157a03c67137022d59393614352d6bf4312"
dependencies = [
- "rand_core 0.10.1",
+ "rand_core 0.10.0",
]
[[package]]
@@ -8898,7 +8906,7 @@ dependencies = [
"num-traits",
"paste",
"profiling",
- "rand 0.9.4",
+ "rand 0.9.2",
"rand_chacha 0.9.0",
"simd_helpers",
"thiserror 2.0.18",
@@ -9253,7 +9261,7 @@ version = "0.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "46a4bd6027df676bcb752d3724db0ea3c0c5fc1dd0376fec51ac7dcaf9cc69be"
dependencies = [
- "rand 0.9.4",
+ "rand 0.9.2",
]
[[package]]
@@ -9636,9 +9644,9 @@ checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f"
[[package]]
name = "rustls-webpki"
-version = "0.103.11"
+version = "0.103.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "20a6af516fea4b20eccceaf166e8aa666ac996208e8a644ce3ef5aa783bc7cd4"
+checksum = "df33b2b81ac578cabaf06b89b0631153a3f416b0a886e8a7a1707fb51abbd1ef"
dependencies = [
"aws-lc-rs",
"ring",
@@ -10261,6 +10269,7 @@ name = "shard"
version = "0.1.0"
dependencies = [
"bytes",
+ "compio",
"consensus",
"crossfire",
"futures",
@@ -10727,9 +10736,9 @@ dependencies = [
[[package]]
name = "sse-stream"
-version = "0.2.2"
+version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2c5e6deb40826033bd7b11c7ef25ef71193fabd71f680f40dd16538a2704d2f4"
+checksum = "eb4dc4d33c68ec1f27d386b5610a351922656e1fdf5c05bbaad930cd1519479a"
dependencies = [
"bytes",
"futures-util",
@@ -11748,7 +11757,7 @@ dependencies = [
"http 1.4.0",
"httparse",
"log",
- "rand 0.9.4",
+ "rand 0.9.2",
"rustls",
"rustls-pki-types",
"sha1",
@@ -11767,7 +11776,7 @@ dependencies = [
"http 1.4.0",
"httparse",
"log",
- "rand 0.9.4",
+ "rand 0.9.2",
"rustls",
"rustls-pki-types",
"sha1",
@@ -11780,7 +11789,7 @@ version = "2.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c"
dependencies = [
- "rand 0.9.4",
+ "rand 0.9.2",
]
[[package]]
@@ -11897,7 +11906,7 @@ version = "1.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "470dbf6591da1b39d43c14523b2b469c86879a53e8b758c8e090a470fe7b1fbe"
dependencies = [
- "rand 0.9.4",
+ "rand 0.9.2",
"web-time",
]
@@ -12322,9 +12331,9 @@ checksum = "b8dad83b4f25e74f184f64c43b150b91efe7647395b42289f38e50566d82855b"
[[package]]
name = "wasm-bindgen"
-version = "0.2.118"
+version = "0.2.117"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0bf938a0bacb0469e83c1e148908bd7d5a6010354cf4fb73279b7447422e3a89"
+checksum = "0551fc1bb415591e3372d0bc4780db7e587d84e2a7e79da121051c5c4b89d0b0"
dependencies = [
"cfg-if",
"once_cell",
@@ -12335,9 +12344,9 @@ dependencies = [
[[package]]
name = "wasm-bindgen-futures"
-version = "0.4.68"
+version = "0.4.67"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f371d383f2fb139252e0bfac3b81b265689bf45b6874af544ffa4c975ac1ebf8"
+checksum = "03623de6905b7206edd0a75f69f747f134b7f0a2323392d664448bf2d3c5d87e"
dependencies = [
"js-sys",
"wasm-bindgen",
@@ -12345,9 +12354,9 @@ dependencies = [
[[package]]
name = "wasm-bindgen-macro"
-version = "0.2.118"
+version = "0.2.117"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "eeff24f84126c0ec2db7a449f0c2ec963c6a49efe0698c4242929da037ca28ed"
+checksum = "7fbdf9a35adf44786aecd5ff89b4563a90325f9da0923236f6104e603c7e86be"
dependencies = [
"quote",
"wasm-bindgen-macro-support",
@@ -12355,9 +12364,9 @@ dependencies = [
[[package]]
name = "wasm-bindgen-macro-support"
-version = "0.2.118"
+version = "0.2.117"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9d08065faf983b2b80a79fd87d8254c409281cf7de75fc4b773019824196c904"
+checksum = "dca9693ef2bab6d4e6707234500350d8dad079eb508dca05530c85dc3a529ff2"
dependencies = [
"bumpalo",
"proc-macro2",
@@ -12368,9 +12377,9 @@ dependencies = [
[[package]]
name = "wasm-bindgen-shared"
-version = "0.2.118"
+version = "0.2.117"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5fd04d9e306f1907bd13c6361b5c6bfc7b3b3c095ed3f8a9246390f8dbdee129"
+checksum = "39129a682a6d2d841b6c429d0c51e5cb0ed1a03829d8b3d1e69a011e62cb3d3b"
dependencies = [
"unicode-ident",
]
@@ -12461,9 +12470,9 @@ dependencies = [
[[package]]
name = "web-sys"
-version = "0.3.95"
+version = "0.3.94"
source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4f2dfbb17949fa2088e5d39408c48368947b86f7834484e87b73de55bc14d97d"
+checksum = "cd70027e39b12f0849461e08ffc50b9cd7688d942c1c8e3c7b22273236b4dd0a"
dependencies = [
"js-sys",
"wasm-bindgen",
diff --git a/DEPENDENCIES.md b/DEPENDENCIES.md
index 243a0fdf3..05575ca2c 100644
--- a/DEPENDENCIES.md
+++ b/DEPENDENCIES.md
@@ -147,7 +147,7 @@ capacity_builder: 0.5.0, "MIT",
capacity_builder_macros: 0.3.0, "MIT",
cargo-platform: 0.3.2, "Apache-2.0 OR MIT",
cargo_metadata: 0.23.1, "MIT",
-cc: 1.2.60, "Apache-2.0 OR MIT",
+cc: 1.2.59, "Apache-2.0 OR MIT",
cesu8: 1.1.0, "Apache-2.0 OR MIT",
cexpr: 0.6.0, "Apache-2.0 OR MIT",
cfg-if: 1.0.4, "Apache-2.0 OR MIT",
@@ -451,7 +451,7 @@ hwlocality-sys: 0.7.0, "MIT",
hybrid-array: 0.4.10, "Apache-2.0 OR MIT",
hyper: 1.9.0, "MIT",
hyper-named-pipe: 0.1.0, "Apache-2.0",
-hyper-rustls: 0.27.8, "Apache-2.0 OR ISC OR MIT",
+hyper-rustls: 0.27.7, "Apache-2.0 OR ISC OR MIT",
hyper-timeout: 0.5.2, "Apache-2.0 OR MIT",
hyper-util: 0.1.20, "MIT",
hyperlocal: 0.9.1, "MIT",
@@ -532,7 +532,7 @@ jni-sys: 0.4.1, "Apache-2.0 OR MIT",
jni-sys-macros: 0.4.1, "Apache-2.0 OR MIT",
jobserver: 0.1.34, "Apache-2.0 OR MIT",
journal: 0.1.0, "Apache-2.0",
-js-sys: 0.3.95, "Apache-2.0 OR MIT",
+js-sys: 0.3.94, "Apache-2.0 OR MIT",
jsonwebtoken: 10.3.0, "MIT",
jwalk: 0.8.1, "MIT",
keccak: 0.2.0, "Apache-2.0 OR MIT",
@@ -654,11 +654,11 @@ once_cell: 1.21.4, "Apache-2.0 OR MIT",
once_cell_polyfill: 1.70.2, "Apache-2.0 OR MIT",
opaque-debug: 0.3.1, "Apache-2.0 OR MIT",
opendal: 0.55.0, "Apache-2.0",
-openssl: 0.10.77, "Apache-2.0",
+openssl: 0.10.76, "Apache-2.0",
openssl-macros: 0.1.1, "Apache-2.0 OR MIT",
openssl-probe: 0.2.1, "Apache-2.0 OR MIT",
openssl-src: 300.6.0+3.6.2, "Apache-2.0 OR MIT",
-openssl-sys: 0.9.113, "MIT",
+openssl-sys: 0.9.112, "MIT",
opentelemetry: 0.31.0, "Apache-2.0",
opentelemetry-appender-tracing: 0.31.1, "Apache-2.0",
opentelemetry-http: 0.31.0, "Apache-2.0",
@@ -708,7 +708,7 @@ pin-utils: 0.1.0, "Apache-2.0 OR MIT",
pinned: 0.1.0, "Apache-2.0 OR MIT",
pkcs1: 0.7.5, "Apache-2.0 OR MIT",
pkcs8: 0.10.2, "Apache-2.0 OR MIT",
-pkg-config: 0.3.33, "Apache-2.0 OR MIT",
+pkg-config: 0.3.32, "Apache-2.0 OR MIT",
plain: 0.2.3, "Apache-2.0 OR MIT",
png: 0.17.16, "Apache-2.0 OR MIT",
png: 0.18.1, "Apache-2.0 OR MIT",
@@ -761,13 +761,13 @@ r-efi: 5.3.0, "Apache-2.0 OR LGPL-2.1-or-later OR MIT",
r-efi: 6.0.0, "Apache-2.0 OR LGPL-2.1-or-later OR MIT",
radium: 0.7.0, "MIT",
rand: 0.8.5, "Apache-2.0 OR MIT",
-rand: 0.9.4, "Apache-2.0 OR MIT",
+rand: 0.9.2, "Apache-2.0 OR MIT",
rand: 0.10.1, "Apache-2.0 OR MIT",
rand_chacha: 0.3.1, "Apache-2.0 OR MIT",
rand_chacha: 0.9.0, "Apache-2.0 OR MIT",
rand_core: 0.6.4, "Apache-2.0 OR MIT",
rand_core: 0.9.5, "Apache-2.0 OR MIT",
-rand_core: 0.10.1, "Apache-2.0 OR MIT",
+rand_core: 0.10.0, "Apache-2.0 OR MIT",
rand_xoshiro: 0.8.0, "Apache-2.0 OR MIT",
rav1e: 0.8.1, "BSD-2-Clause",
ravif: 0.13.0, "BSD-3-Clause",
@@ -826,7 +826,7 @@ rustls-pemfile: 2.2.0, "Apache-2.0 OR ISC OR MIT",
rustls-pki-types: 1.14.0, "Apache-2.0 OR MIT",
rustls-platform-verifier: 0.6.2, "Apache-2.0 OR MIT",
rustls-platform-verifier-android: 0.1.1, "Apache-2.0 OR MIT",
-rustls-webpki: 0.103.11, "ISC",
+rustls-webpki: 0.103.10, "ISC",
rustversion: 1.0.22, "Apache-2.0 OR MIT",
rustybuzz: 0.20.1, "MIT",
ryu: 1.0.23, "Apache-2.0 OR BSL-1.0",
@@ -907,7 +907,7 @@ sqlx-macros-core: 0.8.6, "Apache-2.0 OR MIT",
sqlx-mysql: 0.8.6, "Apache-2.0 OR MIT",
sqlx-postgres: 0.8.6, "Apache-2.0 OR MIT",
sqlx-sqlite: 0.8.6, "Apache-2.0 OR MIT",
-sse-stream: 0.2.2, "Apache-2.0 OR MIT",
+sse-stream: 0.2.1, "Apache-2.0 OR MIT",
stable_deref_trait: 1.2.1, "Apache-2.0 OR MIT",
static-toml: 1.3.0, "MIT",
static_assertions: 1.1.0, "Apache-2.0 OR MIT",
@@ -1063,11 +1063,11 @@ wasi: 0.11.1+wasi-snapshot-preview1, "Apache-2.0 OR Apache-2.0 WITH LLVM-excepti
wasip2: 1.0.2+wasi-0.2.9, "Apache-2.0 OR Apache-2.0 WITH LLVM-exception OR MIT",
wasip3: 0.4.0+wasi-0.3.0-rc-2026-01-06, "Apache-2.0 OR Apache-2.0 WITH LLVM-exception OR MIT",
wasite: 0.1.0, "Apache-2.0 OR BSL-1.0 OR MIT",
-wasm-bindgen: 0.2.118, "Apache-2.0 OR MIT",
-wasm-bindgen-futures: 0.4.68, "Apache-2.0 OR MIT",
-wasm-bindgen-macro: 0.2.118, "Apache-2.0 OR MIT",
-wasm-bindgen-macro-support: 0.2.118, "Apache-2.0 OR MIT",
-wasm-bindgen-shared: 0.2.118, "Apache-2.0 OR MIT",
+wasm-bindgen: 0.2.117, "Apache-2.0 OR MIT",
+wasm-bindgen-futures: 0.4.67, "Apache-2.0 OR MIT",
+wasm-bindgen-macro: 0.2.117, "Apache-2.0 OR MIT",
+wasm-bindgen-macro-support: 0.2.117, "Apache-2.0 OR MIT",
+wasm-bindgen-shared: 0.2.117, "Apache-2.0 OR MIT",
wasm-encoder: 0.244.0, "Apache-2.0 OR Apache-2.0 WITH LLVM-exception OR MIT",
wasm-metadata: 0.244.0, "Apache-2.0 OR Apache-2.0 WITH LLVM-exception OR MIT",
wasm-streams: 0.4.2, "Apache-2.0 OR MIT",
@@ -1075,7 +1075,7 @@ wasm-streams: 0.5.0, "Apache-2.0 OR MIT",
wasm_dep_analyzer: 0.3.0, "MIT",
wasmparser: 0.244.0, "Apache-2.0 OR Apache-2.0 WITH LLVM-exception OR MIT",
wasmtimer: 0.4.3, "MIT",
-web-sys: 0.3.95, "Apache-2.0 OR MIT",
+web-sys: 0.3.94, "Apache-2.0 OR MIT",
web-time: 1.1.0, "Apache-2.0 OR MIT",
webpki-root-certs: 1.0.6, "CDLA-Permissive-2.0",
webpki-roots: 0.26.11, "CDLA-Permissive-2.0",
diff --git a/core/binary_protocol/src/consensus/iobuf.rs b/core/binary_protocol/src/consensus/iobuf.rs
index 4e1355457..2d532d0d3 100644
--- a/core/binary_protocol/src/consensus/iobuf.rs
+++ b/core/binary_protocol/src/consensus/iobuf.rs
@@ -15,14 +15,14 @@
// specific language governing permissions and limitations
// under the License.
+use aligned_vec::{AVec, ConstAlign};
+use compio_buf::{IoBuf, IoBufMut, ReserveError, ReserveExactError, SetLen};
+use std::mem::MaybeUninit;
use std::ops::{Deref, DerefMut, RangeBounds};
use std::ptr::NonNull;
use std::slice;
use std::sync::atomic::{AtomicUsize, Ordering, fence};
-use aligned_vec::{AVec, ConstAlign};
-use compio_buf::IoBuf;
-
#[derive(Debug, Clone)]
pub struct Owned<const ALIGN: usize = 4096> {
inner: AVec<u8, ConstAlign<ALIGN>>,
@@ -91,6 +91,16 @@ impl<const ALIGN: usize> Owned<ALIGN> {
Self { inner }
}
+ /// Allocate a buffer with `capacity` reserved bytes and `len == 0`.
+ /// Intended for use with compio's [`IoBufMut`] read path: the caller
+ /// hands the `Owned` to `read_exact`, which fills it in-place and
+ /// advances the length via [`SetLen::set_len`].
+ pub fn with_capacity(capacity: usize) -> Self {
+ Self {
+ inner: AVec::with_capacity(ALIGN, capacity),
+ }
+ }
+
pub fn copy_from_slice(data: &[u8]) -> Self {
let mut inner: AVec<u8, ConstAlign<ALIGN>> = AVec::new(ALIGN);
inner.extend_from_slice(data);
@@ -133,6 +143,54 @@ impl<const ALIGN: usize> Owned<ALIGN> {
}
}
+// SAFETY contract for the compio IO buf impls below:
+// * `as_init` returns the initialized prefix (len bytes).
+// * `as_uninit` returns the full allocation up to capacity; compio writes
+// bytes into it during reads and THEN calls `SetLen::set_len` to advance
+// the initialized region, which is exactly what `AVec::set_len` expects.
+// * `buf_capacity` is the AVec capacity, not the initialized length.
+impl<const ALIGN: usize> IoBuf for Owned<ALIGN> {
+ fn as_init(&self) -> &[u8] {
+ self.inner.as_slice()
+ }
+}
+
+impl<const ALIGN: usize> IoBufMut for Owned<ALIGN> {
+ fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
+ let cap = self.inner.capacity();
+ let ptr = self.inner.as_mut_ptr().cast::<MaybeUninit<u8>>();
+ unsafe { slice::from_raw_parts_mut(ptr, cap) }
+ }
+
+ fn buf_capacity(&mut self) -> usize {
+ self.inner.capacity()
+ }
+
+ fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
+ // `reserve(additional)` contract: capacity >= current_len + additional
+ // after the call. `AVec::reserve` matches this contract exactly.
+ if self.inner.capacity() - self.inner.len() >= len {
+ return Ok(());
+ }
+ self.inner.reserve(len);
+ Ok(())
+ }
+
+ fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
+ if self.inner.capacity() - self.inner.len() >= len {
+ return Ok(());
+ }
+ self.inner.reserve_exact(len);
+ Ok(())
+ }
+}
+
+impl<const ALIGN: usize> SetLen for Owned<ALIGN> {
+ unsafe fn set_len(&mut self, len: usize) {
+ unsafe { self.inner.set_len(len) };
+ }
+}
+
impl<const ALIGN: usize> Prefix<ALIGN> {
pub fn copy_from_slice(src: &[u8]) -> Self {
Self {
diff --git a/core/binary_protocol/src/consensus/message.rs b/core/binary_protocol/src/consensus/message.rs
index f5b502088..0979511b4 100644
--- a/core/binary_protocol/src/consensus/message.rs
+++ b/core/binary_protocol/src/consensus/message.rs
@@ -24,7 +24,7 @@ use crate::consensus::{
use smallvec::SmallVec;
use std::{marker::PhantomData, mem::size_of};
-const MESSAGE_ALIGN: usize = 4096;
+pub const MESSAGE_ALIGN: usize = 4096;
pub trait MessageBacking<H>
where
diff --git a/core/binary_protocol/src/consensus/mod.rs b/core/binary_protocol/src/consensus/mod.rs
index 0f17f8ab0..5d3ea82a8 100644
--- a/core/binary_protocol/src/consensus/mod.rs
+++ b/core/binary_protocol/src/consensus/mod.rs
@@ -66,7 +66,7 @@ pub use header::{
PrepareOkHeader, ReplyHeader, RequestHeader, StartViewChangeHeader, StartViewHeader,
};
pub use message::{
- ConsensusMessage, FragmentedBacking, Message, MessageBacking, MessageBag, MutableBacking,
- RequestBacking, RequestBackingKind, ResponseBacking, ResponseBackingKind,
+ ConsensusMessage, FragmentedBacking, MESSAGE_ALIGN, Message, MessageBacking, MessageBag,
+ MutableBacking, RequestBacking, RequestBackingKind, ResponseBacking, ResponseBackingKind,
};
pub use operation::Operation;
diff --git a/core/configs/src/server_config/cluster.rs b/core/configs/src/server_config/cluster.rs
index 28421015d..fe2743929 100644
--- a/core/configs/src/server_config/cluster.rs
+++ b/core/configs/src/server_config/cluster.rs
@@ -36,12 +36,28 @@ pub struct NodeConfig {
pub struct CurrentNodeConfig {
pub name: String,
pub ip: String,
+ /// Explicit numeric replica ID for VSR consensus (0-based).
+ ///
+ /// Required when `cluster.enabled` is true. `None` is a configuration
+ /// error caught by [`ClusterConfig::validate`]; a silent default of 0
+ /// would wedge the cluster into a permanent view change as soon as a
+ /// second replica joined.
+ pub replica_id: Option<u8>,
+ /// Bind ports for the current node. Symmetric with [`OtherNodeConfig`]
+ /// so the consensus listener can pick a dedicated `tcp_replica` port
+ /// without borrowing it from the client-facing `tcp.address`.
+ #[serde(default)]
+ pub ports: TransportPorts,
}
#[derive(Debug, Deserialize, Serialize, Clone, ConfigEnv)]
pub struct OtherNodeConfig {
pub name: String,
pub ip: String,
+ /// Explicit numeric replica ID for VSR consensus (0-based).
+ ///
+ /// Required when `cluster.enabled` is true. See [`CurrentNodeConfig::replica_id`].
+ pub replica_id: Option<u8>,
pub ports: TransportPorts,
}
@@ -51,4 +67,6 @@ pub struct TransportPorts {
pub quic: Option<u16>,
pub http: Option<u16>,
pub websocket: Option<u16>,
+ /// Dedicated port for replica-to-replica consensus traffic.
+ pub tcp_replica: Option<u16>,
}
diff --git a/core/configs/src/server_config/defaults.rs b/core/configs/src/server_config/defaults.rs
index 7403196df..1b1d6f6cd 100644
--- a/core/configs/src/server_config/defaults.rs
+++ b/core/configs/src/server_config/defaults.rs
@@ -583,6 +583,8 @@ impl Default for NodeConfig {
current: CurrentNodeConfig {
name: SERVER_CONFIG.cluster.node.current.name.parse().unwrap(),
ip: SERVER_CONFIG.cluster.node.current.ip.parse().unwrap(),
+ replica_id: Some(SERVER_CONFIG.cluster.node.current.replica_id as u8),
+ ports: TransportPorts::default(),
},
others: SERVER_CONFIG
.cluster
@@ -592,11 +594,13 @@ impl Default for NodeConfig {
.map(|other| OtherNodeConfig {
name: other.name.parse().unwrap(),
ip: other.ip.parse().unwrap(),
+ replica_id: Some(other.replica_id as u8),
ports: TransportPorts {
tcp: Some(other.ports.tcp as u16),
quic: Some(other.ports.quic as u16),
http: Some(other.ports.http as u16),
websocket: Some(other.ports.websocket as u16),
+ tcp_replica: Some(other.ports.tcp_replica as u16),
},
})
.collect(),
diff --git a/core/configs/src/server_config/validators.rs b/core/configs/src/server_config/validators.rs
index 6a49157a8..ac491df3c 100644
--- a/core/configs/src/server_config/validators.rs
+++ b/core/configs/src/server_config/validators.rs
@@ -412,6 +412,35 @@ impl Validatable<ConfigurationError> for ClusterConfig {
}
}
+ // VSR needs every replica to have a stable, unique id strictly
+ // less than the total replica count. Missing ids would let serde's
+ // default wedge a backup into thinking it is replica 0, which is
+ // the primary in view 0; duplicate ids split the cluster into two
+ // replicas claiming the same slot; out-of-range ids never win a
+ // primary election. All of these are unrecoverable at runtime -
+ // fail fast at startup.
+ let Some(current_id) = self.node.current.replica_id else {
+ eprintln!(
+ "Invalid cluster configuration: cluster.node.current.replica_id must be set when cluster is enabled"
+ );
+ return Err(ConfigurationError::InvalidConfigurationValue);
+ };
+
+ let total_replicas = u8::try_from(self.node.others.len() + 1).map_err(|_| {
+ eprintln!("Invalid cluster configuration: more than 255 replicas is unsupported");
+ ConfigurationError::InvalidConfigurationValue
+ })?;
+
+ if current_id >= total_replicas {
+ eprintln!(
+ "Invalid cluster configuration: current replica_id {current_id} must be < total replica count {total_replicas}"
+ );
+ return Err(ConfigurationError::InvalidConfigurationValue);
+ }
+
+ let mut seen_ids = std::collections::HashSet::new();
+ seen_ids.insert(current_id);
+
let mut used_endpoints = std::collections::HashSet::new();
for node in &self.node.others {
if node.name.trim().is_empty() {
@@ -427,6 +456,29 @@ impl Validatable<ConfigurationError> for ClusterConfig {
return Err(ConfigurationError::InvalidConfigurationValue);
}
+ let Some(other_id) = node.replica_id else {
+ eprintln!(
+ "Invalid cluster configuration: replica_id must be set for node '{}' when cluster is enabled",
+ node.name
+ );
+ return Err(ConfigurationError::InvalidConfigurationValue);
+ };
+
+ if other_id >= total_replicas {
+ eprintln!(
+ "Invalid cluster configuration: replica_id {other_id} for node '{}' must be < total replica count {total_replicas}",
+ node.name
+ );
+ return Err(ConfigurationError::InvalidConfigurationValue);
+ }
+
+ if !seen_ids.insert(other_id) {
+ eprintln!(
+ "Invalid cluster configuration: duplicate replica_id {other_id} (already held by current node or another other node)"
+ );
+ return Err(ConfigurationError::InvalidConfigurationValue);
+ }
+
let port_list = [
("TCP", node.ports.tcp),
("QUIC", node.ports.quic),
@@ -459,3 +511,74 @@ impl Validatable<ConfigurationError> for ClusterConfig {
Ok(())
}
}
+
+#[cfg(test)]
+mod cluster_validate_tests {
+ use super::*;
+ use crate::server_config::cluster::{
+ ClusterConfig, CurrentNodeConfig, NodeConfig, OtherNodeConfig, TransportPorts,
+ };
+
+ fn other(name: &str, id: Option<u8>) -> OtherNodeConfig {
+ OtherNodeConfig {
+ name: name.to_string(),
+ ip: "127.0.0.1".to_string(),
+ replica_id: id,
+ ports: TransportPorts::default(),
+ }
+ }
+
+ fn cfg(current_id: Option<u8>, others: Vec<OtherNodeConfig>) -> ClusterConfig {
+ ClusterConfig {
+ enabled: true,
+ name: "iggy-cluster".to_string(),
+ node: NodeConfig {
+ current: CurrentNodeConfig {
+ name: "iggy-node-1".to_string(),
+ ip: "127.0.0.1".to_string(),
+ replica_id: current_id,
+ ports: TransportPorts::default(),
+ },
+ others,
+ },
+ }
+ }
+
+ #[test]
+ fn validate_rejects_missing_current_replica_id() {
+ let c = cfg(None, vec![other("n2", Some(1))]);
+ assert!(c.validate().is_err());
+ }
+
+ #[test]
+ fn validate_rejects_missing_other_replica_id() {
+ let c = cfg(Some(0), vec![other("n2", None)]);
+ assert!(c.validate().is_err());
+ }
+
+ #[test]
+ fn validate_rejects_duplicate_replica_ids() {
+ let c = cfg(Some(0), vec![other("n2", Some(0))]);
+ assert!(c.validate().is_err());
+ }
+
+ #[test]
+ fn validate_rejects_out_of_range_replica_id() {
+ // 2 replicas total (current + 1 other), so id 2 is out of range.
+ let c = cfg(Some(2), vec![other("n2", Some(1))]);
+ assert!(c.validate().is_err());
+ }
+
+ #[test]
+ fn validate_accepts_unique_contiguous_replica_ids() {
+ let c = cfg(Some(0), vec![other("n2", Some(1)), other("n3", Some(2))]);
+ assert!(c.validate().is_ok());
+ }
+
+ #[test]
+ fn validate_skips_checks_when_disabled() {
+ let mut c = cfg(None, vec![]);
+ c.enabled = false;
+ assert!(c.validate().is_ok());
+ }
+}
diff --git a/core/consensus/src/impls.rs b/core/consensus/src/impls.rs
index d0932afa0..82e202525 100644
--- a/core/consensus/src/impls.rs
+++ b/core/consensus/src/impls.rs
@@ -1721,16 +1721,20 @@ impl<B: MessageBus, P: Pipeline<Entry = PipelineEntry>> VsrConsensus<B, P> {
#[allow(clippy::future_not_send)]
pub(crate) async fn send_or_loopback(&self, target: u8, message: Message<GenericHeader>)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
{
if target == self.replica {
self.push_loopback(message);
- } else {
- // TODO: Propagate send errors instead of panicking; requires bus error design.
- self.message_bus
- .send_to_replica(target, message)
- .await
- .unwrap();
+ } else if let Err(e) = self
+ .message_bus
+ .send_to_replica(target, message.into_frozen())
+ .await
+ {
+ tracing::warn!(
+ replica = self.replica,
+ target,
+ "send_or_loopback failed: {e}"
+ );
}
}
diff --git a/core/consensus/src/plane_helpers.rs b/core/consensus/src/plane_helpers.rs
index 5a5fecfa9..98e320320 100644
--- a/core/consensus/src/plane_helpers.rs
+++ b/core/consensus/src/plane_helpers.rs
@@ -21,10 +21,8 @@ use crate::{
Status, VsrConsensus,
};
use iggy_binary_protocol::consensus::iobuf::Owned;
-use iggy_binary_protocol::{
- Command2, GenericHeader, Message, PrepareHeader, PrepareOkHeader, ReplyHeader,
-};
-use message_bus::MessageBus;
+use iggy_binary_protocol::{Command2, Message, PrepareHeader, PrepareOkHeader, ReplyHeader};
+use message_bus::{MessageBus, SendError};
use std::ops::AsyncFnOnce;
// TODO: Rework all of those helpers, once the boundaries are more clear and we have a better picture of the commonalities between all of the planes.
@@ -43,7 +41,7 @@ pub async fn request_preflight<B, P>(
request: u64,
) -> Option<Notify>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = PipelineEntry>,
{
let status = consensus
@@ -55,7 +53,7 @@ where
// Best-effort resend, client may have disconnected.
let _ = consensus
.message_bus()
- .send_to_client(client_id, cached_reply.into_generic())
+ .send_to_client(client_id, cached_reply.into_generic().into_frozen())
.await;
None
}
@@ -86,7 +84,7 @@ pub async fn register_preflight<B, P>(
client_id: u128,
) -> Option<Notify>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = PipelineEntry>,
{
let status = consensus.client_table().borrow().check_register(client_id);
@@ -154,7 +152,7 @@ pub const fn fence_old_prepare_by_commit<B, P>(
header: &PrepareHeader,
) -> bool
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = PipelineEntry>,
{
header.op <= consensus.commit_min()
@@ -162,18 +160,25 @@ where
/// Shared chain-replication forwarding to the next replica.
///
+/// Borrows the message, makes a deep copy for the wire, and lets the caller
+/// retain ownership for journal append.
+///
+/// # Errors
+///
+/// Returns `SendError` if the bus fails to deliver to the next replica.
+/// Callers decide error policy (VSR retransmits from WAL via prepare timeout).
+///
/// # Panics
/// - If `header.command` is not `Command2::Prepare`.
/// - If `header.op <= consensus.commit_min()`.
/// - If the computed next replica equals self.
-/// - If the message bus send fails.
#[allow(clippy::future_not_send)]
pub async fn replicate_to_next_in_chain<B, P>(
consensus: &VsrConsensus<B, P>,
- message: Message<PrepareHeader>,
-) -> Message<PrepareHeader>
+ message: &Message<PrepareHeader>,
+) -> Result<(), SendError>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = PipelineEntry>,
{
let header = *message.header();
@@ -185,20 +190,17 @@ where
let primary = consensus.primary_index(header.view);
if next == primary {
- return message;
+ return Ok(());
}
assert_ne!(next, consensus.replica());
- // TODO: Propagate send error instead of panicking; requires bus error design.
- let returned = consensus
- .message_bus()
- .send_to_replica(next, message.into_generic())
- .await
- .unwrap();
- returned
- .try_into_typed()
- .expect("replica send must return the same prepare message")
+ // Chain replication to the next replica is N=1, so the freeze-once
+ // trick does not apply: the caller retains ownership of `message` for
+ // the journal append, so we still deep_copy here. Future refactor
+ // could freeze once and share the backing with the journal path.
+ let frozen = message.deep_copy().into_generic().into_frozen();
+ consensus.message_bus().send_to_replica(next, frozen).await
}
/// Shared preflight checks for `on_replicate`.
@@ -216,7 +218,7 @@ pub fn replicate_preflight<B, P>(
header: &PrepareHeader,
) -> Result<u64, IgnoreReason>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = PipelineEntry>,
{
assert_eq!(header.command, Command2::Prepare);
@@ -249,7 +251,7 @@ where
/// normal status.
pub fn ack_preflight<B, P>(consensus: &VsrConsensus<B, P>) -> Result<(), IgnoreReason>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = PipelineEntry>,
{
if !consensus.is_primary() {
@@ -274,7 +276,7 @@ pub fn ack_quorum_reached<B, P>(
ack: &PrepareOkHeader,
) -> bool
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = PipelineEntry>,
{
if !matches!(
@@ -314,7 +316,7 @@ where
/// If `head()` returns `Some` but `pop()` returns `None` (unreachable).
pub fn drain_committable_prefix<B, P>(consensus: &VsrConsensus<B, P>) -> Vec<PipelineEntry>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = PipelineEntry>,
{
let commit = consensus.commit_max();
@@ -346,7 +348,7 @@ pub fn build_reply_message<B, P>(
body: bytes::Bytes,
) -> Message<ReplyHeader>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = PipelineEntry>,
{
let header_size = std::mem::size_of::<ReplyHeader>();
@@ -417,7 +419,7 @@ pub async fn send_prepare_ok<B, P>(
header: &PrepareHeader,
is_persisted: Option<bool>,
) where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = PipelineEntry>,
{
assert_eq!(header.command, Command2::Prepare);
@@ -476,47 +478,28 @@ pub async fn send_prepare_ok<B, P>(
mod tests {
use super::*;
use crate::{Consensus, LocalPipeline};
- use iggy_common::IggyError;
+ use iggy_binary_protocol::consensus::MESSAGE_ALIGN;
+ use iggy_binary_protocol::consensus::iobuf::Frozen;
+ use message_bus::SendError;
#[derive(Debug, Default)]
struct NoopBus;
impl MessageBus for NoopBus {
- type Client = u128;
- type Replica = u8;
- type Data = Message<GenericHeader>;
- type Sender = ();
-
- fn add_client(&mut self, _client: Self::Client, _sender: Self::Sender) -> bool {
- true
- }
-
- fn remove_client(&mut self, _client: Self::Client) -> bool {
- true
- }
-
- fn add_replica(&mut self, _replica: Self::Replica) -> bool {
- true
- }
-
- fn remove_replica(&mut self, _replica: Self::Replica) -> bool {
- true
- }
-
async fn send_to_client(
&self,
- _client_id: Self::Client,
- data: Self::Data,
- ) -> Result<Self::Data, IggyError> {
- Ok(data)
+ _client_id: u128,
+ _data: Frozen<MESSAGE_ALIGN>,
+ ) -> Result<(), SendError> {
+ Ok(())
}
async fn send_to_replica(
&self,
- _replica: Self::Replica,
- data: Self::Data,
- ) -> Result<Self::Data, IggyError> {
- Ok(data)
+ _replica: u8,
+ _data: Frozen<MESSAGE_ALIGN>,
+ ) -> Result<(), SendError> {
+ Ok(())
}
}
@@ -681,7 +664,7 @@ mod tests {
}
struct SpyBus {
- sent: std::cell::RefCell<Vec<(u8, Message<GenericHeader>)>>,
+ sent: std::cell::RefCell<Vec<(u8, Frozen<MESSAGE_ALIGN>)>>,
}
impl SpyBus {
@@ -694,39 +677,20 @@ mod tests {
#[allow(clippy::future_not_send)]
impl MessageBus for SpyBus {
- type Client = u128;
- type Replica = u8;
- type Data = Message<GenericHeader>;
- type Sender = ();
-
- fn add_client(&mut self, _client: Self::Client, _sender: Self::Sender) -> bool {
- true
- }
- fn remove_client(&mut self, _client: Self::Client) -> bool {
- true
- }
- fn add_replica(&mut self, _replica: Self::Replica) -> bool {
- true
- }
- fn remove_replica(&mut self, _replica: Self::Replica) -> bool {
- true
- }
async fn send_to_client(
&self,
- _client_id: Self::Client,
- data: Self::Data,
- ) -> Result<Self::Data, IggyError> {
- Ok(data)
+ _client_id: u128,
+ _data: Frozen<MESSAGE_ALIGN>,
+ ) -> Result<(), SendError> {
+ Ok(())
}
async fn send_to_replica(
&self,
- replica: Self::Replica,
- data: Self::Data,
- ) -> Result<Self::Data, IggyError> {
- let stored = data.deep_copy();
- let returned = data.deep_copy();
- self.sent.borrow_mut().push((replica, stored));
- Ok(returned)
+ replica: u8,
+ data: Frozen<MESSAGE_ALIGN>,
+ ) -> Result<(), SendError> {
+ self.sent.borrow_mut().push((replica, data));
+ Ok(())
}
}
diff --git a/core/message_bus/Cargo.toml b/core/message_bus/Cargo.toml
index e31049902..cd98b3891 100644
--- a/core/message_bus/Cargo.toml
+++ b/core/message_bus/Cargo.toml
@@ -29,9 +29,16 @@ readme = "../../../README.md"
publish = false
[dependencies]
+async-channel = { workspace = true }
+compio = { workspace = true }
+futures = { workspace = true }
iggy_binary_protocol = { workspace = true }
iggy_common = { workspace = true }
-rand = { workspace = true }
+libc = { workspace = true }
+rand = "0.10"
+socket2 = { workspace = true, features = ["all"] }
+thiserror = { workspace = true }
+tracing = { workspace = true }
[lints.clippy]
enum_glob_use = "deny"
diff --git a/core/message_bus/src/cache/connection.rs b/core/message_bus/src/cache/connection.rs
index a52a260f3..f4373f7d3 100644
--- a/core/message_bus/src/cache/connection.rs
+++ b/core/message_bus/src/cache/connection.rs
@@ -14,18 +14,16 @@
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
+
use crate::cache::AllocationStrategy;
-use iggy_common::TcpSender;
-use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
-use std::{
- cell::RefCell,
- collections::{HashMap, HashSet},
- rc::Rc,
-};
+use rand::SeedableRng;
+use rand::rngs::StdRng;
+use rand::seq::SliceRandom;
+use std::cell::{Cell, RefCell};
+use std::collections::{HashMap, HashSet};
-const MAX_CONNECTIONS_PER_REPLICA: usize = 8;
+pub const DEFAULT_MAX_CONNECTIONS_PER_REPLICA: usize = 8;
-// TODO: Move to some common trait location.
pub trait ShardedState {
type Entry;
type Delta;
@@ -33,30 +31,153 @@ pub trait ShardedState {
fn apply(&mut self, delta: Self::Delta);
}
-/// Least-loaded allocation strategy for connections
+/// Identifies a connection on a specific shard.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct ConnectionAssignment {
+ pub replica: u8,
+ pub shard: u16,
+}
+
+/// Maps a source shard to the shard that owns the connection.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct ShardAssignment {
+ pub replica: u8,
+ pub shard: u16,
+ pub conn_shard: u16,
+}
+
+/// Changeset for connection-based allocation.
+#[derive(Debug, Clone)]
+pub enum ConnectionChanges {
+ Allocate {
+ connections: Vec<ConnectionAssignment>,
+ mappings: Vec<ShardAssignment>,
+ },
+ Deallocate {
+ replica: u8,
+ connections: Vec<ConnectionAssignment>,
+ },
+}
+
+/// Round-robin allocator: assigns each new replica to the next shard in
+/// sequence. Exactly one owning shard per replica.
+///
+/// This is the production strategy today. [`LeastLoadedStrategy`] is kept
+/// alongside as unwired scaffolding for a future tuning pass.
+#[derive(Debug)]
+pub struct RoundRobinStrategy {
+ total_shards: u16,
+ counter: Cell<u16>,
+ assigned: RefCell<HashMap<u8, u16>>,
+}
+
+impl RoundRobinStrategy {
+ #[must_use]
+ pub fn new(total_shards: u16) -> Self {
+ Self {
+ total_shards,
+ counter: Cell::new(0),
+ assigned: RefCell::new(HashMap::new()),
+ }
+ }
+
+ fn next_shard(&self) -> u16 {
+ let next = self.counter.get();
+ self.counter
+ .set(next.wrapping_add(1).rem_euclid(self.total_shards.max(1)));
+ next.rem_euclid(self.total_shards.max(1))
+ }
+}
+
+impl AllocationStrategy<ConnectionCache> for RoundRobinStrategy {
+ fn allocate(&self, replica: u8) -> Option<<ConnectionCache as ShardedState>::Delta> {
+ if self.total_shards == 0 {
+ return None;
+ }
+ if self.assigned.borrow().contains_key(&replica) {
+ return None;
+ }
+
+ let owner = self.next_shard();
+ self.assigned.borrow_mut().insert(replica, owner);
+
+ let connections = vec![ConnectionAssignment {
+ replica,
+ shard: owner,
+ }];
+ let mappings = (0..self.total_shards)
+ .map(|shard| ShardAssignment {
+ replica,
+ shard,
+ conn_shard: owner,
+ })
+ .collect();
+
+ Some(ConnectionChanges::Allocate {
+ connections,
+ mappings,
+ })
+ }
+
+ fn deallocate(&self, replica: u8) -> Option<<ConnectionCache as ShardedState>::Delta> {
+ let owner = self.assigned.borrow_mut().remove(&replica)?;
+ let connections = vec![ConnectionAssignment {
+ replica,
+ shard: owner,
+ }];
+ Some(ConnectionChanges::Deallocate {
+ replica,
+ connections,
+ })
+ }
+}
+
+/// Least-loaded allocation strategy for connections.
+///
+/// NOT wired to production in this revision; see [`RoundRobinStrategy`]. Kept
+/// as scaffolding for a future tuning pass that distributes multiple
+/// connections per replica across the least-loaded shards. The bugs the
+/// earlier version shipped with (non-advancing seed, unfiltered mapping
+/// apply) are fixed here so this code does not rot.
#[derive(Debug)]
pub struct LeastLoadedStrategy {
total_shards: u16,
+ max_connections: usize,
connections_per_shard: RefCell<Vec<(u16, usize)>>,
replica_to_shards: RefCell<HashMap<u8, HashSet<u16>>>,
- rng_seed: u64,
+ rng_seed: Cell<u64>,
}
impl LeastLoadedStrategy {
- pub fn new(total_shards: u16, seed: u64) -> Self {
+ #[must_use]
+ pub fn new(total_shards: u16, max_connections: usize, seed: u64) -> Self {
Self {
total_shards,
+ max_connections,
connections_per_shard: RefCell::new((0..total_shards).map(|s| (s, 0)).collect()),
replica_to_shards: RefCell::new(HashMap::new()),
- rng_seed: seed,
+ rng_seed: Cell::new(seed),
}
}
+ /// Advance the seed via an LCG step so successive allocations see
+ /// different shuffle orders without needing external randomness.
+ fn next_seed(&self) -> u64 {
+ let s = self
+ .rng_seed
+ .get()
+ .wrapping_mul(6_364_136_223_846_793_005)
+ .wrapping_add(1_442_695_040_888_963_407);
+ self.rng_seed.set(s);
+ s
+ }
+
fn create_shard_mappings(
&self,
mappings: &mut Vec<ShardAssignment>,
replica: u8,
mut conn_shards: Vec<u16>,
+ seed: u64,
) {
for &shard in &conn_shards {
mappings.push(ShardAssignment {
@@ -66,7 +187,7 @@ impl LeastLoadedStrategy {
});
}
- let mut rng = StdRng::seed_from_u64(self.rng_seed);
+ let mut rng = StdRng::seed_from_u64(seed);
conn_shards.shuffle(&mut rng);
let mut j = 0;
@@ -83,64 +204,48 @@ impl LeastLoadedStrategy {
j += 1;
}
}
-}
-/// Identifies a connection on a specific shard
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub struct ConnectionAssignment {
- replica: u8,
- shard: u16,
-}
-
-/// Maps a source shard to the shard that owns the connection
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub struct ShardAssignment {
- replica: u8,
- shard: u16,
- conn_shard: u16,
-}
-
-/// Changeset for connection-based allocation
-#[derive(Debug, Clone)]
-pub enum ConnectionChanges {
- Allocate {
- connections: Vec<ConnectionAssignment>,
- mappings: Vec<ShardAssignment>,
- },
- Deallocate {
- connections: Vec<ConnectionAssignment>,
- mappings: Vec<ConnectionAssignment>,
- },
+ /// Returns the set of shard ids that own connections for a given replica.
+ #[must_use]
+ pub fn owning_shards(&self, replica: u8) -> Option<HashSet<u16>> {
+ self.replica_to_shards.borrow().get(&replica).cloned()
+ }
}
-type Delta = <ConnectionCache as ShardedState>::Delta;
impl AllocationStrategy<ConnectionCache> for LeastLoadedStrategy {
- fn allocate(&self, replica: u8) -> Option<Delta> {
+ fn allocate(&self, replica: u8) -> Option<<ConnectionCache as ShardedState>::Delta> {
if self.replica_to_shards.borrow().contains_key(&replica) {
return None;
}
let mut connections = Vec::new();
let mut mappings = Vec::new();
- let connections_needed = usize::from(self.total_shards).min(MAX_CONNECTIONS_PER_REPLICA);
+ let connections_needed = usize::from(self.total_shards).min(self.max_connections);
+ if connections_needed == 0 {
+ return None;
+ }
- let mut rng = StdRng::seed_from_u64(self.rng_seed);
- self.connections_per_shard.borrow_mut().shuffle(&mut rng);
- self.connections_per_shard
- .borrow_mut()
- .sort_by_key(|(_, count)| *count);
+ let seed = self.next_seed();
+ {
+ let mut cps = self.connections_per_shard.borrow_mut();
+ let mut rng = StdRng::seed_from_u64(seed);
+ cps.shuffle(&mut rng);
+ cps.sort_by_key(|(_, count)| *count);
+ }
let mut assigned_shards = HashSet::with_capacity(connections_needed);
- for i in 0..connections_needed {
- let mut connections_per_shard = self.connections_per_shard.borrow_mut();
- let (shard, count) = connections_per_shard.get_mut(i).unwrap();
- connections.push(ConnectionAssignment {
- replica,
- shard: *shard,
- });
- *count += 1;
- assigned_shards.insert(*shard);
+ {
+ let mut cps = self.connections_per_shard.borrow_mut();
+ for i in 0..connections_needed {
+ let (shard, count) = cps.get_mut(i).expect("guarded by connections_needed");
+ connections.push(ConnectionAssignment {
+ replica,
+ shard: *shard,
+ });
+ *count += 1;
+ assigned_shards.insert(*shard);
+ }
}
self.replica_to_shards
@@ -151,20 +256,19 @@ impl AllocationStrategy<ConnectionCache> for LeastLoadedStrategy {
&mut mappings,
replica,
assigned_shards.into_iter().collect(),
+ seed,
);
- Some(Delta::Allocate {
+ Some(ConnectionChanges::Allocate {
connections,
mappings,
})
}
- fn deallocate(&self, replica: u8) -> Option<Delta> {
+ fn deallocate(&self, replica: u8) -> Option<<ConnectionCache as ShardedState>::Delta> {
let conn_shards = self.replica_to_shards.borrow_mut().remove(&replica)?;
let mut connections = Vec::new();
- let mut mappings = Vec::new();
-
for shard in &conn_shards {
if let Some((_, count)) = self
.connections_per_shard
@@ -180,18 +284,14 @@ impl AllocationStrategy<ConnectionCache> for LeastLoadedStrategy {
});
}
- for shard in 0..self.total_shards {
- mappings.push(ConnectionAssignment { replica, shard });
- }
-
- Some(Delta::Deallocate {
+ Some(ConnectionChanges::Deallocate {
+ replica,
connections,
- mappings,
})
}
}
-/// Coordinator that wraps a strategy for a specific sharded state type
+/// Coordinator that wraps a strategy for a specific sharded state type.
#[derive(Debug)]
pub struct Coordinator<A, SS>
where
@@ -240,7 +340,6 @@ where
{
pub fn allocate(&mut self, entry: SS::Entry) -> bool {
if let Some(delta) = self.coordinator.allocate(entry) {
- // TODO: broadcast to other shards.
self.state.apply(delta);
true
} else {
@@ -250,7 +349,6 @@ where
pub fn deallocate(&mut self, entry: SS::Entry) -> bool {
if let Some(delta) = self.coordinator.deallocate(entry) {
- // TODO: broadcast to other shards.
self.state.apply(delta);
true
} else {
@@ -259,26 +357,48 @@ where
}
}
-/// Cache for connection state per shard
+/// Per-shard cache of connection state.
+///
+/// Each shard holds its own `ConnectionCache`. Owning shards have entries in
+/// `owns_connection` for replicas whose TCP connections live on this shard.
+/// All shards (owning or not) have entries in `connection_map` that tell them
+/// which shard owns the connection for each replica.
#[derive(Debug, Default)]
pub struct ConnectionCache {
pub shard_id: u16,
- pub connections: HashMap<u8, Option<Rc<TcpSender>>>,
+ /// Replicas for which THIS shard owns a connection. Value is true once
+ /// the TCP connection has been established (fd received and registered).
+ pub owns_connection: HashMap<u8, bool>,
+ /// For every known replica, the shard that owns the connection.
pub connection_map: HashMap<u8, u16>,
}
impl ConnectionCache {
- pub fn get_connection(&self, replica: u8) -> Option<Rc<TcpSender>> {
- self.connections.get(&replica).and_then(Clone::clone)
+ #[must_use]
+ pub fn is_owner(&self, replica: u8) -> bool {
+ self.owns_connection.contains_key(&replica)
}
- pub fn get_mapped_shard(&self, replica: u8) -> Option<u16> {
+ #[must_use]
+ pub fn owning_shard(&self, replica: u8) -> Option<u16> {
self.connection_map.get(&replica).copied()
}
+
+ pub fn mark_connected(&mut self, replica: u8) {
+ if let Some(connected) = self.owns_connection.get_mut(&replica) {
+ *connected = true;
+ }
+ }
+
+ pub fn mark_disconnected(&mut self, replica: u8) {
+ if let Some(connected) = self.owns_connection.get_mut(&replica) {
+ *connected = false;
+ }
+ }
}
impl ShardedState for ConnectionCache {
- type Entry = u8; // replica id
+ type Entry = u8;
type Delta = ConnectionChanges;
fn apply(&mut self, delta: Self::Delta) {
@@ -289,24 +409,173 @@ impl ShardedState for ConnectionCache {
mappings,
} => {
for conn in connections.iter().filter(|c| c.shard == shard_id) {
- self.connections.insert(conn.replica, None);
+ self.owns_connection.insert(conn.replica, false);
}
- for mapping in &mappings {
+ for mapping in mappings.iter().filter(|m| m.shard == shard_id) {
self.connection_map
.insert(mapping.replica, mapping.conn_shard);
}
}
ConnectionChanges::Deallocate {
+ replica,
connections,
- mappings,
} => {
for conn in connections.iter().filter(|c| c.shard == shard_id) {
- self.connections.remove(&conn.replica);
- }
- for mapping in &mappings {
- self.connection_map.remove(&mapping.replica);
+ self.owns_connection.remove(&conn.replica);
}
+ self.connection_map.remove(&replica);
+ }
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ fn cache(shard_id: u16) -> ConnectionCache {
+ ConnectionCache {
+ shard_id,
+ ..Default::default()
+ }
+ }
+
+ #[test]
+ fn round_robin_allocates_sequentially() {
+ let strategy = RoundRobinStrategy::new(4);
+ let replicas: Vec<u8> = (0..8).collect();
+ let deltas: Vec<_> = replicas.iter().map(|&r| strategy.allocate(r)).collect();
+
+ let owners: Vec<u16> = deltas
+ .iter()
+ .map(|d| match d.as_ref().unwrap() {
+ ConnectionChanges::Allocate { connections, .. } => connections[0].shard,
+ ConnectionChanges::Deallocate { .. } => panic!("expected Allocate"),
+ })
+ .collect();
+ assert_eq!(owners, vec![0, 1, 2, 3, 0, 1, 2, 3]);
+ }
+
+ #[test]
+ fn round_robin_allocate_duplicate_returns_none() {
+ let strategy = RoundRobinStrategy::new(4);
+ let _ = strategy.allocate(1).expect("first allocate");
+ assert!(strategy.allocate(1).is_none());
+ }
+
+ #[test]
+ fn round_robin_deallocate_cleans_up() {
+ let strategy = RoundRobinStrategy::new(4);
+ let mut state = cache(2);
+ let delta = strategy.allocate(5).unwrap();
+ state.apply(delta);
+ let delta = strategy.deallocate(5).unwrap();
+ state.apply(delta);
+ assert!(state.connection_map.is_empty());
+ assert!(state.owns_connection.is_empty());
+ }
+
+ #[test]
+ fn round_robin_zero_shards_is_none() {
+ let strategy = RoundRobinStrategy::new(0);
+ assert!(strategy.allocate(0).is_none());
+ }
+
+ #[test]
+ fn apply_allocate_sets_mapping_on_every_shard_correctly() {
+ // This exercises the shard-filter fix: each shard should see the
+ // correct owning shard for a replica, regardless of how many other
+ // shards were mentioned in the mapping vec.
+ let strategy = RoundRobinStrategy::new(4);
+ let delta = strategy.allocate(7).unwrap();
+
+ for shard_id in 0..4 {
+ let mut state = cache(shard_id);
+ state.apply(delta.clone());
+ assert_eq!(
+ state.connection_map.get(&7).copied(),
+ Some(0),
+ "shard {shard_id} must map replica 7 to owner 0 (rr first pick)"
+ );
+ let should_own = shard_id == 0;
+ assert_eq!(state.is_owner(7), should_own);
+ }
+ }
+
+ #[test]
+ fn apply_allocate_only_owner_marks_owns_connection() {
+ let strategy = RoundRobinStrategy::new(4);
+ // Skip once so next owner is shard 1.
+ let _ = strategy.allocate(100).unwrap();
+ let delta = strategy.allocate(7).unwrap();
+
+ let mut owner = cache(1);
+ let mut non_owner = cache(3);
+ owner.apply(delta.clone());
+ non_owner.apply(delta);
+
+ assert!(owner.is_owner(7));
+ assert_eq!(owner.owning_shard(7), Some(1));
+ assert!(!non_owner.is_owner(7));
+ assert_eq!(non_owner.owning_shard(7), Some(1));
+ }
+
+ #[test]
+ fn least_loaded_allocate_duplicate_returns_none() {
+ let strategy = LeastLoadedStrategy::new(4, 2, 42);
+ let _ = strategy.allocate(1);
+ assert!(strategy.allocate(1).is_none());
+ }
+
+ #[test]
+ fn least_loaded_deallocate_cleans_up() {
+ let strategy = LeastLoadedStrategy::new(4, 2, 42);
+ let mut state = cache(0);
+ let delta = strategy.allocate(1).unwrap();
+ state.apply(delta);
+ assert!(!state.connection_map.is_empty());
+
+ let delta = strategy.deallocate(1).unwrap();
+ state.apply(delta);
+ assert!(state.connection_map.is_empty());
+ assert!(state.owns_connection.is_empty());
+ }
+
+ #[test]
+ fn least_loaded_max_connections_caps_allocation() {
+ let strategy = LeastLoadedStrategy::new(16, 8, 42);
+ let delta = strategy.allocate(1).unwrap();
+ match &delta {
+ ConnectionChanges::Allocate { connections, .. } => {
+ assert_eq!(connections.len(), 8);
}
+ ConnectionChanges::Deallocate { .. } => panic!("expected Allocate"),
}
}
+
+ #[test]
+ fn least_loaded_single_shard_gets_all_connections() {
+ let strategy = LeastLoadedStrategy::new(1, 8, 42);
+ let delta = strategy.allocate(1).unwrap();
+ match &delta {
+ ConnectionChanges::Allocate {
+ connections,
+ mappings,
+ } => {
+ assert_eq!(connections.len(), 1);
+ assert_eq!(mappings.len(), 1);
+ assert_eq!(mappings[0].conn_shard, 0);
+ }
+ ConnectionChanges::Deallocate { .. } => panic!("expected Allocate"),
+ }
+ }
+
+ #[test]
+ fn least_loaded_seed_advances_between_calls() {
+ let strategy = LeastLoadedStrategy::new(4, 2, 42);
+ let first = strategy.rng_seed.get();
+ let _ = strategy.next_seed();
+ let second = strategy.rng_seed.get();
+ assert_ne!(first, second, "seed must advance on next_seed");
+ }
}
diff --git a/core/message_bus/src/cache/mod.rs b/core/message_bus/src/cache/mod.rs
index fd239df97..a3855cd6e 100644
--- a/core/message_bus/src/cache/mod.rs
+++ b/core/message_bus/src/cache/mod.rs
@@ -17,8 +17,7 @@
use crate::cache::connection::ShardedState;
-// TODO: Move to some common trait location.
-/// Allocation strategy that produces deltas for a specific sharded state type
+/// Allocation strategy that produces deltas for a specific sharded state type.
pub trait AllocationStrategy<SS>
where
SS: ShardedState,
diff --git a/core/message_bus/src/client_listener.rs b/core/message_bus/src/client_listener.rs
new file mode 100644
index 000000000..75f13afe5
--- /dev/null
+++ b/core/message_bus/src/client_listener.rs
@@ -0,0 +1,97 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! TCP listener for consensus-protocol SDK clients.
+//!
+//! Runs only on shard 0. Every accepted connection is handed to an
+//! `on_accepted` callback supplied by the shard bootstrap, which mints a
+//! client id (encoding the target shard in the top 16 bits), duplicates
+//! the fd, and ships the setup frame to the owning shard.
+//!
+//! Clients send `RequestHeader` messages and receive `ReplyHeader` messages.
+//! The owning shard's router handler wraps the duplicated fd and installs
+//! writer + reader tasks via [`crate::installer`].
+
+use crate::lifecycle::ShutdownToken;
+use crate::{AcceptedClientFn, GenericHeader, Message};
+use compio::net::{SocketOpts, TcpListener};
+use futures::FutureExt;
+use iggy_common::IggyError;
+use std::net::SocketAddr;
+use std::rc::Rc;
+use tracing::{debug, error, info};
+
+/// Callback preserved for the owning-shard install path in
+/// [`crate::installer`]. Dispatches `(client_id, request_message)` into the
+/// shard's consensus pipeline.
+pub type RequestHandler = Rc<dyn Fn(u128, Message<GenericHeader>)>;
+
+/// Bind the listener and return the bound address without starting the
+/// accept loop. Useful for tests that need the port before driving traffic.
+///
+/// # Errors
+///
+/// Returns [`IggyError::CannotBindToSocket`] if the bind fails.
+#[allow(clippy::future_not_send)]
+pub async fn bind(addr: SocketAddr) -> Result<(TcpListener, SocketAddr), IggyError> {
+ // `SO_REUSEPORT` intentionally not set: only shard 0 binds the client
+ // listener. The shard-0 coordinator round-robins accepts to owning
+ // shards via `ShardFramePayload::ClientConnectionSetup`.
+ let opts = SocketOpts::new().nodelay(true).keepalive(true);
+ let listener = TcpListener::bind_with_options(addr, &opts)
+ .await
+ .map_err(|_| IggyError::CannotBindToSocket(addr.to_string()))?;
+ let actual = listener
+ .local_addr()
+ .map_err(|e| IggyError::IoError(e.to_string()))?;
+ Ok((listener, actual))
+}
+
+/// Run the client listener accept loop until the shutdown token fires. The
+/// `on_accepted` callback owns the accepted stream from that point on.
+///
+/// # Security
+///
+/// This listener has no authentication on accept. Deploy behind an
+/// authenticating reverse proxy or a trusted-network boundary.
+#[allow(clippy::future_not_send)]
+pub async fn run(listener: TcpListener, token: ShutdownToken, on_accepted: AcceptedClientFn) {
+ info!(
+ "Consensus client listener accepting on {:?}",
+ listener.local_addr().ok()
+ );
+
+ loop {
+ futures::select! {
+ () = token.wait().fuse() => {
+ debug!("Consensus client listener shutting down");
+ break;
+ }
+ result = listener.accept().fuse() => {
+ match result {
+ Ok((stream, peer_addr)) => {
+ debug!(%peer_addr, "client accepted, delegating fd");
+ on_accepted(stream);
+ }
+ Err(e) => {
+ error!("Consensus client listener accept failed: {e}");
+ }
+ }
+ }
+ }
+ }
+}
diff --git a/core/message_bus/src/config.rs b/core/message_bus/src/config.rs
new file mode 100644
index 000000000..f1e5bbb76
--- /dev/null
+++ b/core/message_bus/src/config.rs
@@ -0,0 +1,88 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Runtime tunables for the message bus.
+//!
+//! These knobs are consensus-liveness-critical (keepalive + reconnect
+//! timers gate view-change latency; batch sizes gate throughput under
+//! backpressure) and must be deployment-tunable.
+//!
+//! TODO: move this module into `core/configs` (as a `MessageBusConfig`
+//! section on `ServerConfig`) once the downstream bootstrap that
+//! constructs [`crate::IggyMessageBus`] from `ServerConfig` lands.
+//! Keeping the type in-crate for now avoids churning the configs crate
+//! ahead of that wiring.
+
+use std::time::Duration;
+
+/// Aggregated runtime configuration for an `IggyMessageBus` instance.
+///
+/// All fields map onto constants that previously lived inline across
+/// `writer_task`, `framing`, `socket_opts`, `installer`, `connector`,
+/// and `lifecycle::connection_registry`. [`Default`] returns the prior
+/// hardcoded values so existing call sites (tests, single-shard
+/// simulators) can migrate without behaviour change.
+#[derive(Debug, Clone)]
+pub struct MessageBusConfig {
+ /// Maximum number of `BusMessage` entries coalesced into a single
+ /// `writev(2)` call by the writer task. Higher values improve
+ /// syscall amortization at the cost of tail latency.
+ pub max_batch: usize,
+
+ /// Wire-level cap on a single framed message, in bytes. Read-side
+ /// validator; undersize or oversize frames are rejected.
+ pub max_message_size: usize,
+
+ /// Bound on the per-peer mpsc queue. The writer task drains; the
+ /// `send_to_*` path enqueues. Too small drops under burst; too
+ /// large delays backpressure signalling.
+ pub peer_queue_capacity: usize,
+
+ /// Interval between outbound reconnect attempts to peers with
+ /// `peer_id > self_id`.
+ pub reconnect_period: Duration,
+
+ /// TCP keepalive idle timer. See `tcp(7)` `TCP_KEEPIDLE`.
+ pub keepalive_idle: Duration,
+
+ /// TCP keepalive probe interval. See `tcp(7)` `TCP_KEEPINTVL`.
+ pub keepalive_interval: Duration,
+
+ /// TCP keepalive retry count. See `tcp(7)` `TCP_KEEPCNT`. Combined
+ /// with idle + interval this gates the half-open connection
+ /// detection window that VSR view-change timers must accommodate.
+ pub keepalive_retries: u32,
+
+ /// Timeout for per-peer close drain (flush writer, tear down
+ /// reader) before force-cancellation.
+ pub close_peer_timeout: Duration,
+}
+
+impl Default for MessageBusConfig {
+ fn default() -> Self {
+ Self {
+ max_batch: 64,
+ max_message_size: 64 * 1024 * 1024,
+ peer_queue_capacity: 256,
+ reconnect_period: Duration::from_secs(5),
+ keepalive_idle: Duration::from_secs(10),
+ keepalive_interval: Duration::from_secs(5),
+ keepalive_retries: 3,
+ close_peer_timeout: Duration::from_secs(2),
+ }
+ }
+}
diff --git a/core/message_bus/src/connector.rs b/core/message_bus/src/connector.rs
new file mode 100644
index 000000000..641d50683
--- /dev/null
+++ b/core/message_bus/src/connector.rs
@@ -0,0 +1,184 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Outbound replica connector.
+//!
+//! Runs only on shard 0. For each peer replica with `peer_id > self_id` the
+//! connector dials at startup and re-dials on a periodic sweep. On a
+//! successful `Ping` handshake the accepted stream is handed to the
+//! `on_dialed` callback supplied by the shard bootstrap, which duplicates
+//! the fd and ships it to the owning shard via the inter-shard channel
+//! (see `shard::coordinator::ShardZeroCoordinator`).
+
+use crate::IggyMessageBus;
+use crate::framing;
+use crate::lifecycle::ShutdownToken;
+use crate::socket_opts::apply_keepalive_for_connection;
+use crate::{AcceptedReplicaFn, GenericHeader, Message};
+use compio::net::TcpStream;
+use iggy_binary_protocol::{Command2, HEADER_SIZE};
+use std::mem::size_of;
+use std::net::SocketAddr;
+use std::rc::Rc;
+use std::time::Duration;
+use tracing::{debug, info, warn};
+
+/// Default reconnect sweep period.
+///
+/// Equivalent to `MessageBusConfig::default().reconnect_period`; exposed
+/// as a named const for test / bench ergonomics. Kept in sync with the
+/// [`MessageBusConfig::default`] impl. Remove once the configs-crate
+/// migration lands and bootstrap always reads the period from
+/// `ServerConfig`.
+pub const DEFAULT_RECONNECT_PERIOD: Duration = Duration::from_secs(5);
+
+/// Dial every peer with `peer_id > self_id` once, then launch a periodic
+/// sweep in the background. The periodic task handle is tracked on the bus
+/// so graceful shutdown can await it.
+#[allow(clippy::future_not_send)]
+pub async fn start(
+ bus: &Rc<IggyMessageBus>,
+ cluster_id: u128,
+ self_id: u8,
+ peers: Vec<(u8, SocketAddr)>,
+ on_dialed: AcceptedReplicaFn,
+ reconnect_period: Duration,
+) {
+ connect_all(bus, cluster_id, self_id, &peers, &on_dialed).await;
+
+ let handler = on_dialed.clone();
+ let token = bus.token();
+ let bus_for_task = Rc::clone(bus);
+ let handle = compio::runtime::spawn(async move {
+ periodic_reconnect(
+ &bus_for_task,
+ cluster_id,
+ self_id,
+ peers,
+ handler,
+ reconnect_period,
+ token,
+ )
+ .await;
+ });
+ bus.track_background(handle);
+}
+
+#[allow(clippy::future_not_send)]
+async fn connect_all(
+ bus: &Rc<IggyMessageBus>,
+ cluster_id: u128,
+ self_id: u8,
+ peers: &[(u8, SocketAddr)],
+ on_dialed: &AcceptedReplicaFn,
+) {
+ for &(peer_id, addr) in peers {
+ if peer_id <= self_id {
+ continue;
+ }
+ // Skip peers that already have a live mapping on this cluster.
+ // `replicas().contains` covers single-shard deployments where the
+ // connection lives on shard 0; `owning_shard` covers multi-shard
+ // deployments where the fd was delegated to a peer shard but the
+ // mapping broadcast reached shard 0. Either hit means a previous
+ // sweep (or the inbound listener) already installed this peer and
+ // we must not dial again - redialing would tear down the live
+ // socket via the `AlreadyRegistered` race and flap the mapping.
+ if bus.replicas().contains(peer_id) || bus.owning_shard(peer_id).is_some() {
+ debug!(
+ replica = peer_id,
+ "skip reconnect: peer already registered on cluster"
+ );
+ continue;
+ }
+ connect_one(bus, cluster_id, self_id, peer_id, addr, on_dialed).await;
+ }
+}
+
+#[allow(clippy::future_not_send)]
+async fn periodic_reconnect(
+ bus: &Rc<IggyMessageBus>,
+ cluster_id: u128,
+ self_id: u8,
+ peers: Vec<(u8, SocketAddr)>,
+ on_dialed: AcceptedReplicaFn,
+ period: Duration,
+ token: ShutdownToken,
+) {
+ while token.sleep_or_shutdown(period).await {
+ connect_all(bus, cluster_id, self_id, &peers, &on_dialed).await;
+ }
+ debug!("replica reconnect periodic task exiting");
+}
+
+/// Dial a single peer, send the `Ping` handshake, and hand the stream to
+/// `on_dialed` on success. Dial / handshake failures are logged and
+/// swallowed; VSR tolerates missing peers and the periodic sweep retries.
+#[allow(clippy::future_not_send)]
+async fn connect_one(
+ bus: &Rc<IggyMessageBus>,
+ cluster_id: u128,
+ self_id: u8,
+ peer_id: u8,
+ addr: SocketAddr,
+ on_dialed: &AcceptedReplicaFn,
+) {
+ let mut stream = match TcpStream::connect(addr).await {
+ Ok(s) => s,
+ Err(e) => {
+ debug!(replica = peer_id, %addr, "connect failed: {e}");
+ return;
+ }
+ };
+ if let Err(e) = stream.set_nodelay(true) {
+ debug!(replica = peer_id, %addr, "set_nodelay failed: {e}");
+ }
+ let cfg = bus.config();
+ if let Err(e) = apply_keepalive_for_connection(
+ &stream,
+ cfg.keepalive_idle,
+ cfg.keepalive_interval,
+ cfg.keepalive_retries,
+ ) {
+ warn!(
+ replica = peer_id, %addr,
+ "failed to configure TCP keepalive on outbound socket: {e}"
+ );
+ }
+
+ let ping = build_ping_message(cluster_id, self_id);
+ if let Err(e) = framing::write_message(&mut stream, ping).await {
+ warn!(replica = peer_id, %addr, "handshake write failed: {e}");
+ return;
+ }
+
+ info!(replica = peer_id, %addr, "dialed peer replica, delegating fd");
+ on_dialed(stream, peer_id);
+}
+
+/// Build a Ping handshake message identifying our replica id.
+fn build_ping_message(cluster_id: u128, replica_id: u8) -> Message<GenericHeader> {
+ #[allow(clippy::cast_possible_truncation)]
+ Message::<GenericHeader>::new(size_of::<GenericHeader>()).transmute_header(
+ |_, h: &mut GenericHeader| {
+ h.command = Command2::Ping;
+ h.cluster = cluster_id;
+ h.replica = replica_id;
+ h.size = HEADER_SIZE as u32;
+ },
+ )
+}
diff --git a/core/message_bus/src/cache/mod.rs b/core/message_bus/src/error.rs
similarity index 52%
copy from core/message_bus/src/cache/mod.rs
copy to core/message_bus/src/error.rs
index fd239df97..548b325fa 100644
--- a/core/message_bus/src/cache/mod.rs
+++ b/core/message_bus/src/error.rs
@@ -15,16 +15,32 @@
// specific language governing permissions and limitations
// under the License.
-use crate::cache::connection::ShardedState;
-
-// TODO: Move to some common trait location.
-/// Allocation strategy that produces deltas for a specific sharded state type
-pub trait AllocationStrategy<SS>
-where
- SS: ShardedState,
-{
- fn allocate(&self, entry: SS::Entry) -> Option<SS::Delta>;
- fn deallocate(&self, entry: SS::Entry) -> Option<SS::Delta>;
-}
+use thiserror::Error;
+
+/// Transport-level error for `MessageBus::send_to_*` operations.
+///
+/// All variants are non-fatal from the consensus perspective.
+/// VSR handles message loss via timeout-driven retransmission from the WAL.
+#[derive(Debug, Error)]
+pub enum SendError {
+ #[error("client {0} not found")]
+ ClientNotFound(u128),
+
+ #[error("replica {0} not connected")]
+ ReplicaNotConnected(u8),
+
+ #[error("connection closed")]
+ ConnectionClosed,
-pub mod connection;
+ #[error("bus is shutting down")]
+ BusShuttingDown,
+
+ #[error("queue full, message dropped")]
+ Backpressure,
+
+ #[error("inter-shard routing to shard {0} failed")]
+ RoutingFailed(u16),
+
+ #[error("fd duplication failed: errno={0}")]
+ DupFailed(i32),
+}
diff --git a/core/message_bus/src/fd_transfer.rs b/core/message_bus/src/fd_transfer.rs
new file mode 100644
index 000000000..076b47d1f
--- /dev/null
+++ b/core/message_bus/src/fd_transfer.rs
@@ -0,0 +1,125 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! 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
+//! to the same socket. The duplicated fd (a plain `i32`) is sent to
+//! the target shard via the inter-shard channel. The target shard
+//! calls [`wrap_duped_fd`] to construct a compio `TcpStream` on its
+//! own runtime.
+//!
+//! Shard 0 then drops its original `TcpStream`, closing its fd. The
+//! socket stays alive because the duplicated fd still references it
+//! in the kernel's file table.
+
+use compio::net::TcpStream;
+use std::io;
+use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
+use tracing::warn;
+
+/// Duplicate the underlying file descriptor of a TCP stream with
+/// `FD_CLOEXEC` set atomically.
+///
+/// Returns the new fd number. The caller must arrange for the new fd
+/// to be consumed by [`wrap_duped_fd`] on the target shard; leaking
+/// it is a resource leak.
+///
+/// Uses `fcntl(F_DUPFD_CLOEXEC)` rather than `dup(2)` so that any
+/// future `fork`+`exec` child cannot inherit this socket.
+///
+/// # Errors
+///
+/// Returns `io::Error` if the `fcntl(2)` syscall fails.
+pub fn dup_fd(stream: &TcpStream) -> io::Result<RawFd> {
+ let original = stream.as_raw_fd();
+ // SAFETY: F_DUPFD_CLOEXEC allocates the lowest-numbered free fd >= 0
+ // referring to the same open file description as `original`, with
+ // FD_CLOEXEC set atomically. Safe to call on any valid fd.
+ let duped = unsafe { libc::fcntl(original, libc::F_DUPFD_CLOEXEC, 0) };
+ if duped == -1 {
+ return Err(io::Error::last_os_error());
+ }
+ Ok(duped)
+}
+
+/// Wrap a previously duplicated raw fd into a compio `TcpStream`.
+///
+/// Must be called on the target shard's compio runtime so the
+/// `TcpStream` is registered with the correct `io_uring` instance.
+///
+/// # Safety
+///
+/// - `fd` must be a valid, open file descriptor referring to a TCP socket.
+/// - `fd` must not be owned by any other `TcpStream` or resource that
+/// will close it.
+/// - The caller takes ownership: the returned `TcpStream` will close
+/// `fd` on drop.
+#[must_use]
+pub unsafe fn wrap_duped_fd(fd: RawFd) -> TcpStream {
+ unsafe { TcpStream::from_raw_fd(fd) }
+}
+
+/// Close a raw fd without wrapping it in a `TcpStream`.
+///
+/// Used for cleanup on error paths when the fd was duped but the
+/// target shard setup failed. Logs a warning if `close(2)` fails
+/// with anything other than `EINTR`, since that signals a resource
+/// accounting issue (wrong fd, double-close, or kernel fd table
+/// corruption) that callers cannot recover from here.
+pub fn close_fd(fd: RawFd) {
+ // SAFETY: We own this fd and are explicitly closing it.
+ let rc = unsafe { libc::close(fd) };
+ if rc == -1 {
+ let err = io::Error::last_os_error();
+ if err.raw_os_error() != Some(libc::EINTR) {
+ warn!(fd, error = %err, "close(2) failed on duped fd");
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use compio::net::{TcpListener, TcpStream};
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn dup_fd_sets_fd_cloexec() {
+ 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 client = client_res.unwrap();
+
+ let duped = dup_fd(&client).expect("dup_fd failed");
+
+ // SAFETY: F_GETFD is safe on any valid fd; duped was just allocated.
+ let flags = unsafe { libc::fcntl(duped, libc::F_GETFD) };
+ assert!(flags >= 0, "F_GETFD failed: {}", io::Error::last_os_error());
+ assert_ne!(
+ flags & libc::FD_CLOEXEC,
+ 0,
+ "duped fd must have FD_CLOEXEC set"
+ );
+
+ close_fd(duped);
+ }
+}
diff --git a/core/message_bus/src/framing.rs b/core/message_bus/src/framing.rs
new file mode 100644
index 000000000..e7a5502d9
--- /dev/null
+++ b/core/message_bus/src/framing.rs
@@ -0,0 +1,209 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Consensus wire framing for reading/writing `Message<GenericHeader>` over TCP.
+//!
+//! Wire format: `[256-byte header][optional body]`
+//! where `body_len = header.size - HEADER_SIZE`.
+//! Header-only messages have `header.size == HEADER_SIZE` (no body).
+
+use compio::BufResult;
+use compio::buf::{IntoInner, IoBuf, IoBufMut};
+use compio::io::{AsyncReadExt, AsyncWriteExt};
+use iggy_binary_protocol::consensus::MESSAGE_ALIGN;
+use iggy_binary_protocol::consensus::iobuf::Owned;
+use iggy_binary_protocol::{GenericHeader, HEADER_SIZE, Message};
+use iggy_common::IggyError;
+
+/// Default hard ceiling on a single wire frame. Frames above this are
+/// almost certainly a protocol violation or a malicious peer; we drop
+/// the connection rather than try to allocate a massive body buffer.
+///
+/// Equivalent to `MessageBusConfig::default().max_message_size`; kept in
+/// sync with the [`crate::MessageBusConfig::default`] impl. Retained as
+/// a named const for test ergonomics and for callers that have no
+/// `IggyMessageBus` in scope (e.g. standalone handshake helpers).
+pub const MAX_MESSAGE_SIZE: usize = 64 * 1024 * 1024;
+
+// Guards the `header_buf[48..52]` size-field decode below. If `GenericHeader`'s
+// layout ever shifts, this trips at compile time before the runtime read
+// parses garbage.
+const _: () = {
+ assert!(
+ std::mem::offset_of!(GenericHeader, size) == 48,
+ "framing::read_message hardcodes size offset 48; GenericHeader layout changed",
+ );
+};
+
+/// Write a consensus message to a stream. Zero-copy: the message's owned
+/// buffer is handed straight to `io_uring` as a `Frozen`.
+///
+/// Used by the handshake and framing test paths. Hot-path bus traffic goes
+/// through [`crate::writer_task`] which batches many messages into a single
+/// `writev` instead.
+///
+/// # Errors
+///
+/// Returns `IggyError::TcpError` if the write fails.
+#[allow(clippy::future_not_send)] // single-threaded compio
+pub async fn write_message<S: AsyncWriteExt + Unpin>(
+ stream: &mut S,
+ message: Message<GenericHeader>,
+) -> Result<(), IggyError> {
+ let buf = message.into_frozen();
+ stream
+ .write_all(buf)
+ .await
+ .0
+ .map_err(|_| IggyError::TcpError)
+}
+
+/// Read a consensus message from a stream, rejecting frames whose
+/// reported size field exceeds `max_message_size`.
+///
+/// Reads the 256-byte fixed header first, inspects `header.size` to
+/// determine if there's a body, reads the body if present, and
+/// constructs a validated `Message<GenericHeader>`.
+///
+/// Allocation: one `Owned<MESSAGE_ALIGN>` per frame. The header is read
+/// into the first `HEADER_SIZE` bytes, then the body is read into the tail
+/// of the same buffer. No intermediate reassembly copy.
+///
+/// # Errors
+///
+/// Returns `IggyError::ConnectionClosed` on EOF.
+/// Returns `IggyError::TcpError` on I/O errors.
+/// Returns `IggyError::InvalidCommand` if the header fails validation.
+#[allow(clippy::future_not_send)] // single-threaded compio
+pub async fn read_message<S: AsyncReadExt + Unpin>(
+ stream: &mut S,
+ max_message_size: usize,
+) -> Result<Message<GenericHeader>, IggyError> {
+ // Stage 1: allocate one `Owned<MESSAGE_ALIGN>` sized for just the
+ // header. `Owned` implements `IoBufMut`, so compio writes directly
+ // into its backing AVec and advances the length via `SetLen::set_len`.
+ let owned = Owned::<MESSAGE_ALIGN>::with_capacity(HEADER_SIZE);
+ let BufResult(result, owned) = stream.read_exact(owned).await;
+ result.map_err(|e| to_read_error(&e))?;
+
+ let total_size = u32::from_le_bytes(
+ owned.as_slice()[48..52]
+ .try_into()
+ .map_err(|_| IggyError::InvalidCommand)?,
+ ) as usize;
+
+ if !(HEADER_SIZE..=max_message_size).contains(&total_size) {
+ return Err(IggyError::InvalidCommand);
+ }
+
+ if total_size == HEADER_SIZE {
+ return Message::<GenericHeader>::try_from(owned).map_err(|_| IggyError::InvalidCommand);
+ }
+
+ let body_size = total_size - HEADER_SIZE;
+
+ // Stage 2: grow the same `Owned` in place and fill the tail via a
+ // slice read. Total allocations for a body frame: ONE
+ // (`Owned::with_capacity(HEADER_SIZE)` plus one in-place realloc of
+ // the backing AVec). Zero memcpys of the data.
+ let mut owned = owned;
+ let _ = IoBufMut::reserve_exact(&mut owned, body_size);
+ let BufResult(result, slice) = stream
+ .read_exact(owned.slice(HEADER_SIZE..total_size))
+ .await;
+ result.map_err(|e| to_read_error(&e))?;
+ let owned = slice.into_inner();
+
+ Message::<GenericHeader>::try_from(owned).map_err(|_| IggyError::InvalidCommand)
+}
+
+fn to_read_error(e: &std::io::Error) -> IggyError {
+ if e.kind() == std::io::ErrorKind::UnexpectedEof {
+ IggyError::ConnectionClosed
+ } else {
+ IggyError::TcpError
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use compio::net::{TcpListener, TcpStream};
+ use iggy_binary_protocol::Command2;
+
+ #[allow(clippy::cast_possible_truncation)]
+ fn make_header_only(command: Command2) -> Message<GenericHeader> {
+ Message::<GenericHeader>::new(HEADER_SIZE).transmute_header(|_, h: &mut GenericHeader| {
+ h.command = command;
+ h.size = HEADER_SIZE as u32;
+ })
+ }
+
+ #[allow(clippy::future_not_send)] // single-threaded compio
+ async fn local_pair() -> (TcpStream, TcpStream) {
+ 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();
+ (client_res.unwrap(), server)
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)] // single-threaded compio
+ async fn write_then_read_header_only() {
+ let (mut a, mut b) = local_pair().await;
+ let msg = make_header_only(Command2::Ping);
+ write_message(&mut a, msg).await.unwrap();
+ let read = read_message(&mut b, MAX_MESSAGE_SIZE).await.unwrap();
+ assert_eq!(read.header().command, Command2::Ping);
+ assert_eq!(read.header().size as usize, HEADER_SIZE);
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)] // single-threaded compio
+ async fn read_rejects_oversize_size_field() {
+ use compio::io::AsyncWriteExt;
+ let (mut a, mut b) = local_pair().await;
+
+ let mut buf = vec![0u8; HEADER_SIZE];
+ let bogus = u32::try_from(MAX_MESSAGE_SIZE + 1)
+ .unwrap_or(u32::MAX)
+ .to_le_bytes();
+ buf[48..52].copy_from_slice(&bogus);
+ a.write_all(buf).await.0.unwrap();
+
+ let res = read_message(&mut b, MAX_MESSAGE_SIZE).await;
+ assert!(matches!(res, Err(IggyError::InvalidCommand)));
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)] // single-threaded compio
+ async fn read_rejects_undersize_size_field() {
+ use compio::io::AsyncWriteExt;
+ let (mut a, mut b) = local_pair().await;
+
+ let mut buf = vec![0u8; HEADER_SIZE];
+ let undersize = u32::try_from(HEADER_SIZE - 1).unwrap();
+ buf[48..52].copy_from_slice(&undersize.to_le_bytes());
+ a.write_all(buf).await.0.unwrap();
+
+ let res = read_message(&mut b, MAX_MESSAGE_SIZE).await;
+ assert!(matches!(res, Err(IggyError::InvalidCommand)));
+ }
+}
diff --git a/core/message_bus/src/installer.rs b/core/message_bus/src/installer.rs
new file mode 100644
index 000000000..af70d92f7
--- /dev/null
+++ b/core/message_bus/src/installer.rs
@@ -0,0 +1,396 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Connection installer trait.
+//!
+//! Shard 0 accepts / dials all TCP connections and ships the duplicated fd
+//! to the owning shard via the inter-shard `ShardFrame` channel. The owning
+//! shard's router handler wraps the fd on its own compio runtime and
+//! registers the connection on its local bus. This trait exposes that
+//! registration surface in a way the shard layer can call without knowing
+//! the concrete bus type.
+
+use crate::client_listener::RequestHandler;
+use crate::fd_transfer;
+use crate::framing;
+use crate::replica_listener::MessageHandler;
+use crate::socket_opts::apply_keepalive_for_connection;
+use crate::{IggyMessageBus, lifecycle::ShutdownToken};
+use compio::net::{OwnedReadHalf, TcpStream};
+use futures::FutureExt;
+use iggy_binary_protocol::Command2;
+use std::cell::Cell;
+use std::os::unix::io::RawFd;
+use std::rc::Rc;
+use tracing::{debug, info, warn};
+
+/// Operations a shard needs to perform on its local bus when the router
+/// receives an inter-shard connection-setup or mapping 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`.
+pub trait ConnectionInstaller {
+ /// Wrap a duplicated TCP fd into a `TcpStream` on the local compio
+ /// runtime, spawn writer + reader tasks, and register the replica
+ /// connection on this shard.
+ ///
+ /// Takes ownership of `fd`. On registration failure the fd is closed
+ /// via `close_fd` (by dropping the wrapping `TcpStream`).
+ fn install_replica_fd(&self, fd: RawFd, replica_id: u8, on_message: MessageHandler);
+
+ /// Same for an SDK client connection. The owning shard is already
+ /// encoded in the top 16 bits of `client_id`.
+ fn install_client_fd(&self, fd: RawFd, client_id: u128, on_request: RequestHandler);
+
+ /// Update the replica -> owning shard mapping used by the `send_to_replica`
+ /// slow path on non-owning shards.
+ fn set_shard_mapping(&self, replica: u8, owning_shard: u16);
+
+ /// Forget the replica -> owning shard mapping (e.g. after a connection
+ /// loss, before the next allocate).
+ fn remove_shard_mapping(&self, replica: u8);
+}
+
+impl ConnectionInstaller for Rc<IggyMessageBus> {
+ fn install_replica_fd(&self, fd: RawFd, replica_id: u8, on_message: MessageHandler) {
+ // SAFETY: Shard 0 duped the fd via `dup(2)` and sent it to this
+ // shard via the inter-shard channel. Ownership is being transferred
+ // to us; no other TcpStream holds it.
+ let stream = unsafe { fd_transfer::wrap_duped_fd(fd) };
+ install_replica_stream(self, replica_id, stream, on_message);
+ }
+
+ fn install_client_fd(&self, fd: RawFd, client_id: u128, on_request: RequestHandler) {
+ // SAFETY: see install_replica_fd.
+ let stream = unsafe { fd_transfer::wrap_duped_fd(fd) };
+ install_client_stream(self, client_id, stream, on_request);
+ }
+
+ fn set_shard_mapping(&self, replica: u8, owning_shard: u16) {
+ IggyMessageBus::set_shard_mapping(self, replica, owning_shard);
+ }
+
+ fn remove_shard_mapping(&self, replica: u8) {
+ IggyMessageBus::remove_shard_mapping(self, replica);
+ }
+}
+
+/// Install a pre-wrapped replica TCP stream on the bus. Shared by
+/// `install_replica_fd` and, once shard 0 acts as the delegate target for
+/// itself, the accept path.
+#[allow(clippy::future_not_send)] // single-threaded compio
+pub fn install_replica_stream(
+ bus: &Rc<IggyMessageBus>,
+ peer_id: u8,
+ stream: TcpStream,
+ on_message: MessageHandler,
+) {
+ if bus.replicas().contains(peer_id) {
+ debug!(
+ replica = peer_id,
+ "replica already registered on this shard, dropping delegated fd"
+ );
+ drop(stream);
+ return;
+ }
+
+ let cfg = bus.config();
+ if let Err(e) = apply_keepalive_for_connection(
+ &stream,
+ cfg.keepalive_idle,
+ cfg.keepalive_interval,
+ cfg.keepalive_retries,
+ ) {
+ warn!(replica = peer_id, "keepalive failed on delegated fd: {e}");
+ }
+
+ let (read_half, write_half) = stream.into_split();
+ let (tx, rx) = async_channel::bounded(bus.peer_queue_capacity());
+
+ // Writer and reader both observe abnormal close and used to fire
+ // `notify_connection_lost` twice per disconnect, causing shard 0 to
+ // broadcast two `ReplicaMappingClear` rounds and churn the mapping.
+ // Shared one-shot guard: whichever half exits first wins.
+ let notified = Rc::new(Cell::new(false));
+ // If the registry insert below races with a concurrent install for
+ // the same peer id and loses, both spawned halves must skip their
+ // post-loop cleanup: calling `replicas().remove` / `close_peer` or
+ // `notify_connection_lost` would evict the winner's entry and
+ // clobber its mapping broadcast. `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));
+
+ let writer_token = bus.token();
+ let bus_for_writer = Rc::clone(bus);
+ let writer_label = format!("{peer_id}");
+ let notified_writer = Rc::clone(¬ified);
+ let aborted_writer = Rc::clone(&install_aborted);
+ let max_batch = bus.config().max_batch;
+ let writer_handle = compio::runtime::spawn(async move {
+ crate::writer_task::run(
+ rx,
+ write_half,
+ writer_token,
+ "replica",
+ writer_label,
+ max_batch,
+ )
+ .await;
+ if aborted_writer.get() || bus_for_writer.is_shutting_down() {
+ return;
+ }
+ bus_for_writer.replicas().remove(peer_id);
+ if !notified_writer.replace(true) {
+ bus_for_writer.notify_connection_lost(peer_id);
+ }
+ });
+
+ let bus_for_reader = Rc::clone(bus);
+ let read_token = bus.token();
+ let notified_reader = Rc::clone(¬ified);
+ let aborted_reader = Rc::clone(&install_aborted);
+ let close_peer_timeout = bus.config().close_peer_timeout;
+ let max_message_size = bus.config().max_message_size;
+ let reader_handle = compio::runtime::spawn(async move {
+ replica_read_loop(
+ peer_id,
+ read_half,
+ &on_message,
+ &read_token,
+ max_message_size,
+ )
+ .await;
+ if aborted_reader.get() {
+ debug!(
+ replica = peer_id,
+ "aborted replica install: skipping post-loop cleanup"
+ );
+ return;
+ }
+ if !read_token.is_triggered() {
+ bus_for_reader
+ .replicas()
+ .close_peer(peer_id, close_peer_timeout)
+ .await;
+ if !notified_reader.replace(true) {
+ bus_for_reader.notify_connection_lost(peer_id);
+ }
+ }
+ info!(replica = peer_id, "peer replica disconnected");
+ });
+
+ if bus
+ .replicas()
+ .insert(peer_id, tx, writer_handle, reader_handle)
+ .is_err()
+ {
+ // Tell both halves to stand down: the winner's entry is live and
+ // must not be touched by this losing install. Dropping the
+ // `async_channel::Sender` (which `insert` already did on error)
+ // will wake the writer with `Err(Closed)`; the writer then drops
+ // its `OwnedWriteHalf`, the peer sees EOF, and the reader's
+ // `read_message` returns with an error. Both halves exit and hit
+ // their `install_aborted` check.
+ install_aborted.set(true);
+ warn!(replica = peer_id, "replica registry insert raced");
+ }
+}
+
+/// Install a pre-wrapped client TCP stream on the bus.
+#[allow(clippy::future_not_send)] // single-threaded compio
+pub fn install_client_stream(
+ bus: &Rc<IggyMessageBus>,
+ client_id: u128,
+ stream: TcpStream,
+ on_request: RequestHandler,
+) {
+ let cfg = bus.config();
+ if let Err(e) = apply_keepalive_for_connection(
+ &stream,
+ cfg.keepalive_idle,
+ cfg.keepalive_interval,
+ cfg.keepalive_retries,
+ ) {
+ warn!(
+ client = client_id,
+ "keepalive failed on delegated client fd: {e}"
+ );
+ }
+
+ let (read_half, write_half) = stream.into_split();
+ let (tx, rx) = async_channel::bounded(bus.peer_queue_capacity());
+
+ // If the registry insert below loses a race for `client_id`, the
+ // losing reader must NOT invoke `on_request` (it would route
+ // responses through the wrong registry entry) and must NOT call
+ // `close_peer` (it would evict the winner). See the replica path
+ // above for the same pattern.
+ let install_aborted = Rc::new(Cell::new(false));
+
+ let writer_token = bus.token();
+ let bus_for_writer = Rc::clone(bus);
+ let writer_label = format!("{client_id:#034x}");
+ let aborted_writer = Rc::clone(&install_aborted);
+ let max_batch = bus.config().max_batch;
+ let writer_handle = compio::runtime::spawn(async move {
+ crate::writer_task::run(
+ rx,
+ write_half,
+ writer_token,
+ "client",
+ writer_label,
+ max_batch,
+ )
+ .await;
+ if aborted_writer.get() || bus_for_writer.is_shutting_down() {
+ return;
+ }
+ bus_for_writer.clients().remove(client_id);
+ });
+
+ let bus_for_reader = Rc::clone(bus);
+ let read_token = bus.token();
+ let aborted_reader = Rc::clone(&install_aborted);
+ let close_peer_timeout = bus.config().close_peer_timeout;
+ let max_message_size = bus.config().max_message_size;
+ let reader_handle = compio::runtime::spawn(async move {
+ client_read_loop(
+ client_id,
+ read_half,
+ &on_request,
+ &read_token,
+ &aborted_reader,
+ max_message_size,
+ )
+ .await;
+ if aborted_reader.get() {
+ debug!(
+ client = client_id,
+ "aborted client install: skipping post-loop cleanup"
+ );
+ return;
+ }
+ if !read_token.is_triggered() {
+ bus_for_reader
+ .clients()
+ .close_peer(client_id, close_peer_timeout)
+ .await;
+ }
+ info!(client = client_id, "consensus client disconnected");
+ });
+
+ if bus
+ .clients()
+ .insert(client_id, tx, writer_handle, reader_handle)
+ .is_err()
+ {
+ // Shard 0 mints client ids as `(target_shard << 112) | seq` with a
+ // monotonic `seq` starting at 1, so wrap requires 2^112 mints and
+ // a collision here is a bootstrap bug or a foreign id leaking
+ // into the setup path. Flip `install_aborted` so the orphan
+ // reader drops inbound frames instead of forwarding them via
+ // `on_request` (which would route responses through the winner's
+ // entry and silently misroute).
+ install_aborted.set(true);
+ warn!(
+ client_id,
+ "duplicate client id in registry, dropping delegated fd \
+ (shard 0 counter invariant violated)"
+ );
+ }
+}
+
+/// Read loop for a delegated replica connection. Identical to the
+/// `replica_listener` version but kept here to avoid cross-module coupling.
+#[allow(clippy::future_not_send)]
+async fn replica_read_loop(
+ replica_id: u8,
+ mut read_half: OwnedReadHalf<TcpStream>,
+ on_message: &MessageHandler,
+ token: &ShutdownToken,
+ max_message_size: usize,
+) {
+ loop {
+ futures::select! {
+ () = token.wait().fuse() => {
+ debug!(replica = replica_id, "replica read loop shutting down");
+ return;
+ }
+ result = framing::read_message(&mut read_half, max_message_size).fuse() => {
+ match result {
+ Ok(msg) => on_message(replica_id, msg),
+ Err(e) => {
+ debug!(replica = replica_id, "read error: {e}");
+ return;
+ }
+ }
+ }
+ }
+ }
+}
+
+/// Read loop for a delegated client connection. Rejects any command other
+/// than `Request` (the client side of the consensus protocol only speaks
+/// request/reply).
+///
+/// `aborted` is set by the installer when the registry insert loses a
+/// duplicate-client-id race. The loop checks it before dispatching each
+/// request so the losing reader can never invoke `on_request` with the
+/// client id owned by the winning install.
+#[allow(clippy::future_not_send)]
+async fn client_read_loop(
+ client_id: u128,
+ mut read_half: OwnedReadHalf<TcpStream>,
+ on_request: &RequestHandler,
+ token: &ShutdownToken,
+ aborted: &Cell<bool>,
+ max_message_size: usize,
+) {
+ loop {
+ futures::select! {
+ () = token.wait().fuse() => {
+ debug!(client = client_id, "client read loop shutting down");
+ return;
+ }
+ result = framing::read_message(&mut read_half, max_message_size).fuse() => {
+ match result {
+ Ok(msg) => {
+ if aborted.get() {
+ return;
+ }
+ let cmd = msg.header().command;
+ if cmd != Command2::Request {
+ warn!(
+ client = client_id,
+ ?cmd,
+ "unexpected command from client, expected Request"
+ );
+ continue;
+ }
+ on_request(client_id, msg);
+ }
+ Err(e) => {
+ debug!(client = client_id, "client read error: {e}");
+ return;
+ }
+ }
+ }
+ }
+ }
+}
diff --git a/core/message_bus/src/lib.rs b/core/message_bus/src/lib.rs
index 5b96b4bcf..de78547f0 100644
--- a/core/message_bus/src/lib.rs
+++ b/core/message_bus/src/lib.rs
@@ -14,124 +14,604 @@
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
-mod cache;
-use crate::cache::connection::{
- ConnectionCache, Coordinator, LeastLoadedStrategy, ShardedConnections,
+pub mod cache;
+pub mod client_listener;
+pub mod config;
+pub mod connector;
+mod error;
+pub mod fd_transfer;
+pub mod framing;
+pub mod installer;
+pub mod lifecycle;
+pub mod replica_io;
+pub mod replica_listener;
+pub(crate) mod socket_opts;
+pub mod transports;
+pub mod writer_task;
+
+pub use config::MessageBusConfig;
+pub use error::SendError;
+pub use installer::ConnectionInstaller;
+pub use lifecycle::{
+ AlreadyRegistered, BusMessage, BusReceiver, BusSender, ConnectionRegistry, DrainOutcome,
+ ReplicaRegistry, Shutdown, ShutdownToken,
};
+
+use compio::runtime::JoinHandle;
+use iggy_binary_protocol::consensus::MESSAGE_ALIGN;
+use iggy_binary_protocol::consensus::iobuf::Frozen;
use iggy_binary_protocol::{GenericHeader, Message};
-use iggy_common::{IggyError, SenderKind, TcpSender};
-use std::{collections::HashMap, rc::Rc};
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::time::Duration;
-/// Message bus parameterized by allocation strategy and sharded state
-pub trait MessageBus {
- type Client;
- type Replica;
- type Data;
- type Sender;
+/// Callback for forwarding a consensus message to a remote shard.
+///
+/// Provided by the shard layer at bus construction and installed via
+/// [`IggyMessageBus::set_replica_forward_fn`] / [`IggyMessageBus::set_client_forward_fn`].
+/// Arguments: `(target_shard_id, message)`. Returns `Ok(())` on successful
+/// enqueue into the inter-shard channel, or [`SendError::RoutingFailed`] on
+/// `try_send` failure.
+///
+/// Asymmetry vs the `Rc<dyn Fn>` siblings (`AcceptedReplicaFn`,
+/// `AcceptedClientFn`, `ConnectionLostFn`): the forward fn has a single
+/// owner (the bus) and is installed once at bootstrap; `Box` suffices.
+/// The `Rc` siblings are shared with accept loops and connection tasks
+/// that outlive the caller and need independent ownership.
+pub type ShardForwardFn = Box<dyn Fn(u16, Frozen<MESSAGE_ALIGN>) -> Result<(), SendError>>;
+
+/// Callback invoked on every successful replica handshake.
+///
+/// Fired by the replica listener / outbound connector. The callback decides
+/// whether to install the stream locally or ship it to another shard; this
+/// crate does not need to care about that policy. Takes ownership of the
+/// `TcpStream`; the replica id has already been validated against the
+/// cluster config.
+pub type AcceptedReplicaFn = std::rc::Rc<dyn Fn(compio::net::TcpStream, u8)>;
- fn add_client(&mut self, client: Self::Client, sender: Self::Sender) -> bool;
- fn remove_client(&mut self, client: Self::Client) -> bool;
+/// Callback invoked on every accepted SDK client connection.
+///
+/// Takes ownership of the accepted stream and is responsible for minting /
+/// assigning the client id as part of its delegation policy.
+pub type AcceptedClientFn = std::rc::Rc<dyn Fn(compio::net::TcpStream)>;
- fn add_replica(&mut self, replica: Self::Replica) -> bool;
- fn remove_replica(&mut self, replica: Self::Replica) -> bool;
+/// Notifier fired when a delegated replica connection dies.
+///
+/// The delegated replica connection's writer / reader tasks invoke this on
+/// abnormal exit (peer closed, write failed). The shard bootstrap wraps
+/// this closure around a `try_send` into shard 0's inbox so shard 0 can
+/// clear the replica mapping and re-dial.
+pub type ConnectionLostFn = std::rc::Rc<dyn Fn(u8)>;
- // TODO: refactor consesus headers.
+/// Point-to-point message delivery between consensus participants.
+///
+/// `Ok(())` means "accepted for delivery" - NOT "delivered to peer".
+/// The bus never retries. Consensus owns retransmission via the WAL
+/// (Prepare) or VSR timeouts (view change).
+///
+/// The implementation is fire-and-forget: `send_to_*` enqueues the message
+/// to a per-peer bounded mpsc and returns immediately. A dedicated writer
+/// task per connection drains the queue and writes batched frames in a
+/// single `writev` syscall. A slow peer cannot stall sends to other peers.
+///
+/// # Yield semantics
+///
+/// For the production `IggyMessageBus` impl, `send_to_*` has zero `.await`
+/// points in its body; the `async fn` shell exists solely for trait
+/// compatibility with simulator implementations. Callers can assume the
+/// send completes synchronously and the returned future is always `Ready`
+/// on first poll. Future transports (QUIC, TLS) that would introduce real
+/// yields must advertise that change in their own doc; consensus code
+/// relies on the no-yield property to reason about reentrancy.
+pub trait MessageBus {
fn send_to_client(
&self,
- client_id: Self::Client,
- data: Self::Data,
- ) -> impl Future<Output = Result<Self::Data, IggyError>>;
+ client_id: u128,
+ data: Frozen<MESSAGE_ALIGN>,
+ ) -> impl Future<Output = Result<(), SendError>>;
+
fn send_to_replica(
&self,
- replica: Self::Replica,
- data: Self::Data,
- ) -> impl Future<Output = Result<Self::Data, IggyError>>;
+ replica: u8,
+ data: Frozen<MESSAGE_ALIGN>,
+ ) -> impl Future<Output = Result<(), SendError>>;
}
-// TODO: explore generics for Strategy
-#[derive(Debug)]
+/// Production message bus backed by real TCP connections.
+///
+/// Owns:
+/// - a root [`Shutdown`] / [`ShutdownToken`] that fans cancellation to every
+/// accept loop, read loop, periodic task, and per-connection writer task,
+/// - a [`ConnectionRegistry<u128>`] for clients (`u128` id is shard-packed),
+/// - a [`ReplicaRegistry`] for replicas (`u8` id from the Ping handshake,
+/// backed by a fixed-size array to avoid hash lookup on every send),
+/// - the `JoinHandle`s of background tasks (accept loops, reconnect periodic).
+///
+/// Send semantics:
+/// - `send_to_*` clones the per-peer `Sender` out of the registry, calls
+/// `try_send` on it, and returns. No `.await` happens in the body. Drops
+/// on `Full` (returned as [`SendError::Backpressure`]) are recovered by
+/// VSR retransmission.
+/// - The per-connection writer task batches up to `config.max_batch` messages into
+/// a single `writev` syscall via [`writer_task::run`].
+///
+/// Interior mutability via `RefCell` / `Cell` is sound because compio is a
+/// single-threaded runtime: no other task can execute while we hold a borrow.
pub struct IggyMessageBus {
- clients: HashMap<u128, SenderKind>,
- replicas: ShardedConnections<LeastLoadedStrategy, ConnectionCache>,
shard_id: u16,
+ shutdown: Shutdown,
+ token: ShutdownToken,
+ clients: ConnectionRegistry<u128>,
+ replicas: ReplicaRegistry,
+ background_tasks: RefCell<Vec<JoinHandle<()>>>,
+ config: MessageBusConfig,
+ /// For each replica, the shard that owns the TCP connection from this
+ /// shard's perspective. Present on ALL shards (owning and non-owning).
+ /// On the owning shard, `shard_mapping[replica] == self.shard_id`.
+ shard_mapping: RefCell<HashMap<u8, u16>>,
+ /// Forwards a replica-bound message to the shard that owns the replica's
+ /// TCP connection. `None` on single-shard setups and tests.
+ replica_forward_fn: RefCell<Option<ShardForwardFn>>,
+ /// Forwards a client-bound message to the shard that owns the client's
+ /// TCP connection (the owning shard is encoded in the top 16 bits of
+ /// the client id). `None` on single-shard setups and tests.
+ client_forward_fn: RefCell<Option<ShardForwardFn>>,
+ /// Invoked by a delegated replica connection's writer / reader tasks
+ /// when they exit abnormally. `None` when running without a shard-0
+ /// coordinator (single-shard deployments and tests).
+ connection_lost_fn: RefCell<Option<ConnectionLostFn>>,
}
impl IggyMessageBus {
+ /// Construct a bus with [`MessageBusConfig::default`] tunables.
#[must_use]
- pub fn new(total_shards: u16, shard_id: u16, seed: u64) -> Self {
+ pub fn new(shard_id: u16) -> Self {
+ Self::with_config(shard_id, MessageBusConfig::default())
+ }
+
+ /// Construct a bus with a custom per-peer queue capacity; all other
+ /// tunables fall back to [`MessageBusConfig::default`].
+ ///
+ /// Tuning knob for tests and benchmarks. Production should use
+ /// [`with_config`](Self::with_config).
+ #[must_use]
+ pub fn with_capacity(shard_id: u16, peer_queue_capacity: usize) -> Self {
+ let cfg = MessageBusConfig {
+ peer_queue_capacity,
+ ..MessageBusConfig::default()
+ };
+ Self::with_config(shard_id, cfg)
+ }
+
+ /// Construct a bus with the given [`MessageBusConfig`].
+ #[must_use]
+ pub fn with_config(shard_id: u16, config: MessageBusConfig) -> Self {
+ let (shutdown, token) = Shutdown::new();
Self {
- clients: HashMap::new(),
- replicas: ShardedConnections {
- coordinator: Coordinator::new(LeastLoadedStrategy::new(total_shards, seed)),
- state: ConnectionCache {
- shard_id,
- ..Default::default()
- },
- },
shard_id,
+ shutdown,
+ token,
+ clients: ConnectionRegistry::new(),
+ replicas: ReplicaRegistry::new(),
+ background_tasks: RefCell::new(Vec::new()),
+ config,
+ shard_mapping: RefCell::new(HashMap::new()),
+ replica_forward_fn: RefCell::new(None),
+ client_forward_fn: RefCell::new(None),
+ connection_lost_fn: RefCell::new(None),
}
}
- pub fn get_replica_connection(&self, replica: u8) -> Option<Rc<TcpSender>> {
- let mapped_shard = self.replicas.state.get_mapped_shard(replica)?;
- if mapped_shard == self.shard_id {
- self.replicas.state.get_connection(replica)
- } else {
- None
- }
+ /// Install the notifier used by delegated replica connections to tell
+ /// shard 0 that a connection died. Single place to inject in tests too.
+ pub fn set_connection_lost_fn(&self, f: ConnectionLostFn) {
+ *self.connection_lost_fn.borrow_mut() = Some(f);
}
-}
-#[allow(clippy::future_not_send)] // Single-threaded runtime (compio), Rc/RefCell by design
-impl MessageBus for IggyMessageBus {
- type Client = u128;
- type Replica = u8;
- type Data = Message<GenericHeader>;
- type Sender = SenderKind;
-
- fn add_client(&mut self, client: Self::Client, sender: Self::Sender) -> bool {
- if self.clients.contains_key(&client) {
- return false;
+ /// Invoke the registered connection-lost notifier, if any.
+ pub(crate) fn notify_connection_lost(&self, replica_id: u8) {
+ if let Some(f) = self.connection_lost_fn.borrow().as_ref() {
+ f(replica_id);
}
- self.clients.insert(client, sender);
- true
}
- fn remove_client(&mut self, client: Self::Client) -> bool {
- self.clients.remove(&client).is_some()
+ /// Install the replica-plane inter-shard forward closure.
+ ///
+ /// Non-owning shards invoke this from `send_to_replica`'s slow path to
+ /// push the message into the owning shard's inbox. The owning shard's
+ /// router then re-enters `send_to_replica` on the local bus (fast path).
+ ///
+ /// Takes `&self` so it can be called through an `Rc<IggyMessageBus>`
+ /// wrapper after the bus is shared with accept loops and periodic tasks.
+ /// Must not be called from within a forward closure invocation
+ /// (re-entrant `borrow_mut` panics) — call sites are bootstrap only.
+ pub fn set_replica_forward_fn(&self, f: ShardForwardFn) {
+ *self.replica_forward_fn.borrow_mut() = Some(f);
+ }
+
+ /// Install the client-plane inter-shard forward closure.
+ ///
+ /// Shards invoke this from `send_to_client`'s slow path when the client
+ /// connection lives on a different shard (top 16 bits of `client_id`).
+ ///
+ /// Takes `&self` so it can be called through an `Rc<IggyMessageBus>`
+ /// wrapper. Must not be called from within a forward closure
+ /// invocation. Call sites are bootstrap only.
+ pub fn set_client_forward_fn(&self, f: ShardForwardFn) {
+ *self.client_forward_fn.borrow_mut() = Some(f);
+ }
+
+ /// Update the shard mapping for a replica.
+ ///
+ /// Called when the allocation strategy assigns or reassigns connections.
+ pub fn set_shard_mapping(&self, replica: u8, owning_shard: u16) {
+ self.shard_mapping
+ .borrow_mut()
+ .insert(replica, owning_shard);
}
- fn add_replica(&mut self, replica: Self::Replica) -> bool {
- self.replicas.allocate(replica)
+ /// Remove all shard mappings for a replica (on deallocation).
+ pub fn remove_shard_mapping(&self, replica: u8) {
+ self.shard_mapping.borrow_mut().remove(&replica);
}
- fn remove_replica(&mut self, replica: Self::Replica) -> bool {
- self.replicas.deallocate(replica)
+ /// Get the owning shard for a replica, if mapped.
+ pub fn owning_shard(&self, replica: u8) -> Option<u16> {
+ self.shard_mapping.borrow().get(&replica).copied()
+ }
+
+ #[must_use]
+ pub const fn shard_id(&self) -> u16 {
+ self.shard_id
+ }
+
+ /// Per-peer mpsc capacity used when registering new connections.
+ #[must_use]
+ pub const fn peer_queue_capacity(&self) -> usize {
+ self.config.peer_queue_capacity
}
+ /// Runtime tunables in effect on this bus.
+ #[must_use]
+ pub const fn config(&self) -> &MessageBusConfig {
+ &self.config
+ }
+
+ /// Cheap clone of the root shutdown token.
+ ///
+ /// Handed to accept loops, read tasks, writer tasks, and periodic tasks
+ /// so they can `select!` on cancellation.
+ #[must_use]
+ pub fn token(&self) -> ShutdownToken {
+ self.token.clone()
+ }
+
+ /// Whether [`shutdown`](Self::shutdown) has been called.
+ #[must_use]
+ pub fn is_shutting_down(&self) -> bool {
+ self.shutdown.is_triggered()
+ }
+
+ /// Accessor used by the client listener to insert / remove connections
+ /// and by `send_to_client` to look up senders.
+ #[must_use]
+ pub const fn clients(&self) -> &ConnectionRegistry<u128> {
+ &self.clients
+ }
+
+ /// Accessor used by the replica listener and connector to insert /
+ /// remove peer connections and by `send_to_replica` to look up senders.
+ #[must_use]
+ pub const fn replicas(&self) -> &ReplicaRegistry {
+ &self.replicas
+ }
+
+ /// Register a background task (accept loop, reconnect periodic) so
+ /// [`shutdown`](Self::shutdown) can await it.
+ ///
+ /// The tracking vec grows during shutdown too. `shutdown` drains it in
+ /// a loop until empty, so a task pushed mid-shutdown is still awaited.
+ pub fn track_background(&self, handle: JoinHandle<()>) {
+ self.background_tasks.borrow_mut().push(handle);
+ }
+
+ /// Trigger the root shutdown and drain everything with the given
+ /// deadline.
+ ///
+ /// Order:
+ /// 1. Trigger the root shutdown token (every accept loop, read loop,
+ /// writer loop, and periodic task selecting on it observes the
+ /// cancellation and exits).
+ /// 2. Drain the client registry (closes each per-peer `Sender` then
+ /// awaits both writer + reader handles).
+ /// 3. Drain the replica registry.
+ /// 4. Loop-drain every tracked background task. Tasks pushed
+ /// mid-shutdown (e.g. a reader that observed the token and
+ /// spawned its own cleanup) are picked up on the next iteration.
+ ///
+ /// Connections drain before background tasks so that writer tasks
+ /// get the full remaining budget for `write_vectored_all` before any
+ /// accept / reconnect / refresh periodic consumes it. Background
+ /// tasks hold no in-flight wire frames, so force-cancelling them
+ /// cannot truncate a frame on the wire.
+ #[allow(clippy::future_not_send)] // single-threaded compio
+ pub async fn shutdown(&self, timeout: Duration) -> DrainOutcome {
+ self.shutdown.trigger();
+
+ let deadline = std::time::Instant::now() + timeout;
+
+ let remaining = deadline.saturating_duration_since(std::time::Instant::now());
+ let clients_outcome = self.clients.drain(remaining).await;
+ let remaining = deadline.saturating_duration_since(std::time::Instant::now());
+ let replicas_outcome = self.replicas.drain(remaining).await;
+
+ let mut background_clean = 0usize;
+ let mut background_force = 0usize;
+ loop {
+ let batch: Vec<JoinHandle<()>> = self.background_tasks.borrow_mut().drain(..).collect();
+ if batch.is_empty() {
+ break;
+ }
+ for handle in batch {
+ let remaining = deadline.saturating_duration_since(std::time::Instant::now());
+ if remaining.is_zero() {
+ drop(handle);
+ background_force += 1;
+ continue;
+ }
+ match compio::time::timeout(remaining, handle).await {
+ Ok(_) => background_clean += 1,
+ Err(_) => background_force += 1,
+ }
+ }
+ }
+
+ DrainOutcome {
+ clean: clients_outcome.clean + replicas_outcome.clean,
+ force: clients_outcome.force + replicas_outcome.force,
+ background_clean,
+ background_force,
+ }
+ }
+}
+
+/// Forwarding impl so `VsrConsensus<Rc<IggyMessageBus>>` (two consensus
+/// planes sharing one bus) type-checks without duplicating the bus or
+/// taking it by value.
+#[allow(clippy::future_not_send)] // single-threaded compio
+impl<T: MessageBus + ?Sized> MessageBus for std::rc::Rc<T> {
+ fn send_to_client(
+ &self,
+ client_id: u128,
+ data: Frozen<MESSAGE_ALIGN>,
+ ) -> impl Future<Output = Result<(), SendError>> {
+ (**self).send_to_client(client_id, data)
+ }
+
+ fn send_to_replica(
+ &self,
+ replica: u8,
+ data: Frozen<MESSAGE_ALIGN>,
+ ) -> impl Future<Output = Result<(), SendError>> {
+ (**self).send_to_replica(replica, data)
+ }
+}
+
+#[allow(clippy::future_not_send)] // single-threaded compio
+impl MessageBus for IggyMessageBus {
async fn send_to_client(
&self,
- client_id: Self::Client,
- message: Self::Data,
- ) -> Result<Self::Data, IggyError> {
- #[allow(clippy::cast_possible_truncation)] // IggyError::ClientNotFound takes u32
- let _sender = self
- .clients
- .get(&client_id)
- .ok_or(IggyError::ClientNotFound(client_id as u32))?;
- Ok(message)
+ client_id: u128,
+ message: Frozen<MESSAGE_ALIGN>,
+ ) -> Result<(), SendError> {
+ if self.is_shutting_down() {
+ return Err(SendError::BusShuttingDown);
+ }
+ // Owning shard is encoded in the top 16 bits of client_id.
+ let owning_shard = client_id_owning_shard(client_id);
+ if owning_shard == self.shard_id {
+ let sender = self
+ .clients
+ .borrow_sender(client_id)
+ .ok_or(SendError::ClientNotFound(client_id))?;
+ return sender.try_send(message).map_err(map_try_send_err);
+ }
+ let forward = self.client_forward_fn.borrow();
+ let forward = forward
+ .as_ref()
+ .ok_or(SendError::ClientNotFound(client_id))?;
+ forward(owning_shard, message)
}
async fn send_to_replica(
&self,
- replica: Self::Replica,
- message: Self::Data,
- ) -> Result<Self::Data, IggyError> {
- // TODO: Handle lazily creating the connection.
- let _connection = self
- .get_replica_connection(replica)
- .ok_or(IggyError::ResourceNotFound(format!("Replica {replica}")))?;
- Ok(message)
+ replica: u8,
+ message: Frozen<MESSAGE_ALIGN>,
+ ) -> Result<(), SendError> {
+ if self.is_shutting_down() {
+ return Err(SendError::BusShuttingDown);
+ }
+ // Fast path: this shard owns a connection to the replica.
+ if let Some(sender) = self.replicas.borrow_sender(replica) {
+ return sender.try_send(message).map_err(map_try_send_err);
+ }
+ // Slow path: route via the inter-shard channel to the owning shard.
+ let owning_shard = self
+ .shard_mapping
+ .borrow()
+ .get(&replica)
+ .copied()
+ .ok_or(SendError::ReplicaNotConnected(replica))?;
+ let forward = self.replica_forward_fn.borrow();
+ let forward = forward
+ .as_ref()
+ .ok_or(SendError::ReplicaNotConnected(replica))?;
+ forward(owning_shard, message)
+ }
+}
+
+/// Extract the owning shard from a client id.
+///
+/// Shard 0 mints client ids as `(target_shard_id << 112) | seq`. The top 16
+/// bits encode which shard's bus registry holds the connection; any shard
+/// that needs to reply to this client uses this accessor to decide between
+/// the fast path (local) and the slow path (forward via inter-shard).
+#[must_use]
+#[allow(clippy::cast_possible_truncation)]
+pub const fn client_id_owning_shard(client_id: u128) -> u16 {
+ (client_id >> 112) as u16
+}
+
+/// Map an `async_channel::TrySendError` onto the bus-level [`SendError`].
+///
+/// Shape-matches `Result::map_err` (takes the error by value) so it can be
+/// used directly as a function reference rather than a closure.
+#[allow(clippy::needless_pass_by_value)] // signature required by map_err
+fn map_try_send_err(e: async_channel::TrySendError<Frozen<MESSAGE_ALIGN>>) -> SendError {
+ match e {
+ async_channel::TrySendError::Full(_) => SendError::Backpressure,
+ async_channel::TrySendError::Closed(_) => SendError::ConnectionClosed,
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use iggy_binary_protocol::{Command2, HEADER_SIZE};
+ use std::cell::RefCell;
+
+ #[allow(clippy::cast_possible_truncation)]
+ fn dummy_message() -> Message<GenericHeader> {
+ Message::<GenericHeader>::new(HEADER_SIZE).transmute_header(|_, h: &mut GenericHeader| {
+ h.command = Command2::Prepare;
+ h.size = HEADER_SIZE as u32;
+ })
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)] // single-threaded compio
+ async fn send_to_client_slow_path_forwards_to_owning_shard() {
+ // Bus on shard 5; client id encodes owning shard = 7.
+ let bus = IggyMessageBus::new(5);
+ let captured: std::rc::Rc<RefCell<Option<u16>>> = std::rc::Rc::new(RefCell::new(None));
+ let captured_clone = captured.clone();
+ bus.set_client_forward_fn(Box::new(move |target, _msg| {
+ *captured_clone.borrow_mut() = Some(target);
+ Ok(())
+ }));
+
+ let client_id = (7u128 << 112) | 0x2a;
+ bus.send_to_client(client_id, dummy_message().into_frozen())
+ .await
+ .expect("forward_fn should accept");
+ assert_eq!(*captured.borrow(), Some(7));
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn send_to_client_fast_path_hits_local_when_owning_shard_matches() {
+ // shard_id == top-16-bits => fast path, registry miss => ClientNotFound
+ let bus = IggyMessageBus::new(3);
+ let client_id = (3u128 << 112) | 1;
+ let err = bus
+ .send_to_client(client_id, dummy_message().into_frozen())
+ .await
+ .unwrap_err();
+ assert!(matches!(err, SendError::ClientNotFound(_)));
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn send_to_client_returns_not_found_when_remote_and_no_forward_fn() {
+ // shard_id != top-16-bits and no forward_fn installed.
+ let bus = IggyMessageBus::new(0);
+ let client_id = (9u128 << 112) | 1;
+ let err = bus
+ .send_to_client(client_id, dummy_message().into_frozen())
+ .await
+ .unwrap_err();
+ assert!(matches!(err, SendError::ClientNotFound(_)));
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn send_to_replica_slow_path_uses_shard_mapping_and_forwards() {
+ let bus = IggyMessageBus::new(0);
+ let captured: std::rc::Rc<RefCell<Option<u16>>> = std::rc::Rc::new(RefCell::new(None));
+ let captured_clone = captured.clone();
+ bus.set_replica_forward_fn(Box::new(move |target, _msg| {
+ *captured_clone.borrow_mut() = Some(target);
+ Ok(())
+ }));
+ bus.set_shard_mapping(5, 3);
+
+ bus.send_to_replica(5, dummy_message().into_frozen())
+ .await
+ .expect("forward ok");
+
+ assert_eq!(*captured.borrow(), Some(3));
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn send_to_replica_no_mapping_returns_not_connected() {
+ let bus = IggyMessageBus::new(0);
+ let err = bus
+ .send_to_replica(9, dummy_message().into_frozen())
+ .await
+ .unwrap_err();
+ assert!(matches!(err, SendError::ReplicaNotConnected(9)));
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn client_id_owning_shard_extracts_top_16_bits() {
+ assert_eq!(client_id_owning_shard((7u128 << 112) | 0x2a), 7);
+ assert_eq!(client_id_owning_shard(0), 0);
+ assert_eq!(
+ client_id_owning_shard((u128::from(u16::MAX)) << 112),
+ u16::MAX
+ );
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)] // single-threaded compio
+ async fn track_background_queues_handles() {
+ let bus = IggyMessageBus::new(0);
+
+ let h1 = compio::runtime::spawn(async {});
+ let h2 = compio::runtime::spawn(async {});
+ bus.track_background(h1);
+ bus.track_background(h2);
+ assert_eq!(bus.background_tasks.borrow().len(), 2);
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)] // single-threaded compio
+ async fn shutdown_loop_drains_tasks_added_during_shutdown() {
+ // Models the real race: a background task spawned after shutdown
+ // has been triggered but before the loop-drain catches it up. The
+ // loop must pick the newer handle up on its next iteration.
+ let bus = IggyMessageBus::new(0);
+
+ let h1 = compio::runtime::spawn(async {});
+ bus.track_background(h1);
+
+ // Simulate mid-shutdown push: trigger, then push a fresh handle
+ // via the public API (imitating a reader task that observed the
+ // token and registered a cleanup future).
+ bus.shutdown.trigger();
+ let h2 = compio::runtime::spawn(async {});
+ bus.track_background(h2);
+
+ let outcome = bus.shutdown(Duration::from_secs(2)).await;
+ assert_eq!(outcome.background_clean, 2);
+ assert_eq!(outcome.background_force, 0);
+ assert_eq!(
+ bus.background_tasks.borrow().len(),
+ 0,
+ "shutdown must leave background_tasks empty",
+ );
}
}
diff --git a/core/message_bus/src/lifecycle/connection_registry.rs b/core/message_bus/src/lifecycle/connection_registry.rs
new file mode 100644
index 000000000..53f1f7da1
--- /dev/null
+++ b/core/message_bus/src/lifecycle/connection_registry.rs
@@ -0,0 +1,682 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Keyed connection registry shared by the client and replica TCP paths.
+//!
+//! Each entry owns:
+//! - the `BusSender` used by `MessageBus::send_to_*` (a `try_send` into the
+//! per-peer bounded mpsc),
+//! - the `JoinHandle` of the writer task (drains the mpsc and pushes batched
+//! `writev` to the wire),
+//! - the `JoinHandle` of the reader task (read loop that hands inbound
+//! messages to the consumer's sync callback).
+//!
+//! Coordination is via the bus-wide [`ShutdownToken`]: triggering it makes
+//! reader/writer tasks observe cancellation and exit. `drain` additionally
+//! closes each `Sender` so writer tasks see the channel close and finish
+//! any in-flight batch before exiting.
+//!
+//! [`ShutdownToken`]: crate::lifecycle::ShutdownToken
+
+use compio::runtime::JoinHandle;
+use iggy_binary_protocol::consensus::MESSAGE_ALIGN;
+use iggy_binary_protocol::consensus::iobuf::Frozen;
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::fmt::Debug;
+use std::hash::Hash;
+use std::time::{Duration, Instant};
+use thiserror::Error;
+use tracing::{debug, warn};
+
+/// Minimum deadline reserved for reader-handle drain, in addition to the
+/// writer drain's remaining budget.
+///
+/// A slow writer that consumes the shared deadline used to leave zero
+/// runway for the reader handle, force-cancelling it even though reader
+/// exit is usually immediate once the socket half-closes. Reserving a
+/// small floor keeps the reader drain observable (`Clean` vs `Force`)
+/// without meaningfully extending total shutdown time.
+pub const READER_DRAIN_FLOOR: Duration = Duration::from_millis(250);
+
+/// Payload type carried over every per-peer queue.
+///
+/// Consensus messages are `Frozen<MESSAGE_ALIGN>` by the time they hit
+/// this queue: the dispatch layer freezes once and fan-out becomes a
+/// refcount bump per target. The writer task reads `Frozen` out of the
+/// queue and passes it straight to `write_vectored_all`, so no
+/// conversion happens on the hot path.
+pub type BusMessage = Frozen<MESSAGE_ALIGN>;
+
+/// Producer side of a per-peer queue. Cloned out of the registry by
+/// `send_to_*` and used with `try_send`.
+pub type BusSender = async_channel::Sender<BusMessage>;
+
+/// Consumer side of a per-peer queue. Owned by the writer task.
+pub type BusReceiver = async_channel::Receiver<BusMessage>;
+
+/// Attempt to register a duplicate key.
+#[derive(Debug, Error)]
+#[error("key already registered")]
+pub struct AlreadyRegistered;
+
+/// Result of [`ConnectionRegistry::drain`] or [`crate::IggyMessageBus::shutdown`].
+///
+/// Counts are aggregate across all drained entries. `background_*` apply
+/// only to bus-level background tasks (accept loops, reconnect periodic)
+/// and stay zero when returned from [`ConnectionRegistry::drain`] directly.
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
+pub struct DrainOutcome {
+ /// Writer / reader tasks that exited within the deadline.
+ pub clean: usize,
+ /// Writer / reader tasks that had to be cancelled when the deadline
+ /// elapsed.
+ pub force: usize,
+ /// Background tasks (accept loops, reconnect periodic) that exited
+ /// within the deadline.
+ pub background_clean: usize,
+ /// Background tasks that had to be cancelled when the deadline
+ /// elapsed.
+ pub background_force: usize,
+}
+
+#[derive(Debug)]
+struct Entry {
+ sender: BusSender,
+ writer_handle: Option<JoinHandle<()>>,
+ reader_handle: Option<JoinHandle<()>>,
+}
+
+/// Map of live connections keyed by some transport-specific id.
+///
+/// For clients `K = u128` (the minted client id); for replicas `K = u8`
+/// (the replica id carried in the Ping handshake).
+#[derive(Debug)]
+pub struct ConnectionRegistry<K>
+where
+ K: Eq + Hash + Copy + Debug + 'static,
+{
+ entries: RefCell<HashMap<K, Entry>>,
+}
+
+impl<K> Default for ConnectionRegistry<K>
+where
+ K: Eq + Hash + Copy + Debug + 'static,
+{
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl<K> ConnectionRegistry<K>
+where
+ K: Eq + Hash + Copy + Debug + 'static,
+{
+ #[must_use]
+ pub fn new() -> Self {
+ Self {
+ entries: RefCell::new(HashMap::new()),
+ }
+ }
+
+ #[must_use]
+ pub fn contains(&self, key: K) -> bool {
+ self.entries.borrow().contains_key(&key)
+ }
+
+ #[must_use]
+ pub fn len(&self) -> usize {
+ self.entries.borrow().len()
+ }
+
+ #[must_use]
+ pub fn is_empty(&self) -> bool {
+ self.entries.borrow().is_empty()
+ }
+
+ /// Borrow the per-peer queue producer without cloning the `Sender`.
+ ///
+ /// `async_channel::Sender::try_send` takes `&self` so a borrow is
+ /// sufficient; cloning would trigger an atomic RMW on the inner
+ /// `Arc<State>` on every send.
+ ///
+ /// The returned guard holds a `RefCell` read borrow. Callers must
+ /// not invoke a mutating method on the same registry
+ /// (`insert` / `remove` / `close_peer`) while the guard is live -
+ /// that would panic on the `borrow_mut`. All `try_send` paths are
+ /// sync so no yield can occur while the guard is held.
+ pub fn borrow_sender(&self, key: K) -> Option<std::cell::Ref<'_, BusSender>> {
+ std::cell::Ref::filter_map(self.entries.borrow(), |entries| {
+ entries.get(&key).map(|entry| &entry.sender)
+ })
+ .ok()
+ }
+
+ /// Register a new connection. Stores the producer side of the per-peer
+ /// queue alongside the writer + reader task handles so graceful shutdown
+ /// can drain everything.
+ ///
+ /// # Errors
+ ///
+ /// Returns [`AlreadyRegistered`] if `key` is already present. Callers
+ /// resolve races (e.g. duplicate dial) by dropping the duplicate.
+ pub fn insert(
+ &self,
+ key: K,
+ sender: BusSender,
+ writer_handle: JoinHandle<()>,
+ reader_handle: JoinHandle<()>,
+ ) -> Result<(), AlreadyRegistered> {
+ let mut entries = self.entries.borrow_mut();
+ if entries.contains_key(&key) {
+ return Err(AlreadyRegistered);
+ }
+ entries.insert(
+ key,
+ Entry {
+ sender,
+ writer_handle: Some(writer_handle),
+ reader_handle: Some(reader_handle),
+ },
+ );
+ Ok(())
+ }
+
+ /// Unregister a connection without awaiting its tasks.
+ ///
+ /// Returns `true` if an entry was removed. Dropping the entry drops the
+ /// `Sender`, which makes the writer task observe `Closed` and exit. The
+ /// reader task is independent and will exit on its own (read error) or
+ /// when the bus shutdown token fires.
+ ///
+ /// Prefer [`close_peer`](Self::close_peer) when closing from the reader
+ /// task: the explicit close-sender, await-writer sequence prevents a
+ /// mid-writev cancellation from landing a truncated frame on the wire.
+ pub fn remove(&self, key: K) -> bool {
+ self.entries.borrow_mut().remove(&key).is_some()
+ }
+
+ /// Close the entry keyed by `key` in the correct order: remove the
+ /// entry from the registry, close its `BusSender` (so the writer task
+ /// observes `Closed` and drains any in-flight batch cleanly), then
+ /// await the writer handle up to `timeout`.
+ ///
+ /// The reader handle is dropped without awaiting because the typical
+ /// caller of `close_peer` IS the reader task (self-remove path).
+ /// Awaiting your own `JoinHandle` would deadlock.
+ ///
+ /// Returns a best-effort result: this is a lifecycle convenience, not
+ /// a signal of a problem if the writer needed to be force-cancelled
+ /// at the deadline.
+ #[allow(clippy::future_not_send)] // single-threaded compio
+ pub async fn close_peer(&self, key: K, timeout: Duration) {
+ let Some(mut entry) = self.entries.borrow_mut().remove(&key) else {
+ return;
+ };
+ entry.sender.close();
+ if let Some(writer_handle) = entry.writer_handle.take() {
+ let _ = compio::time::timeout(timeout, writer_handle).await;
+ }
+ drop(entry.reader_handle);
+ }
+
+ /// Drain every entry, awaiting each task with a shared deadline.
+ ///
+ /// Order:
+ /// 1. Take the entries out of the map (so concurrent inserts during
+ /// drain go to a fresh registry state - we do not block them).
+ /// 2. Close each `Sender` so the writer task sees the channel close.
+ /// Combined with the bus shutdown token, both reader and writer
+ /// will exit cleanly.
+ /// 3. Await both task handles per entry against the shared deadline.
+ /// A handle that does not finish in time is force-cancelled by drop.
+ ///
+ /// Entries are drained concurrently via a `FuturesUnordered`: total
+ /// drain time is bounded by the slowest entry, not the sum across
+ /// entries. Writer-before-reader sequencing is preserved inside each
+ /// entry's future so a mid-writev cancellation cannot truncate a frame.
+ #[allow(clippy::future_not_send)] // single-threaded compio
+ pub async fn drain(&self, timeout: Duration) -> DrainOutcome {
+ let drained: Vec<(K, Entry)> = self.entries.borrow_mut().drain().collect();
+ drain_entries(drained, timeout).await
+ }
+}
+
+/// Shared parallel-drain routine used by both [`ConnectionRegistry`] and
+/// [`ReplicaRegistry`].
+///
+/// Callers take their entries out of storage in whatever way suits their
+/// backing type and hand the pre-collected `Vec` to this helper. Each
+/// entry's sender is closed, writer awaited, then reader awaited, all
+/// concurrently across entries via `FuturesUnordered`.
+#[allow(clippy::future_not_send)] // single-threaded compio
+async fn drain_entries<K>(drained: Vec<(K, Entry)>, timeout: Duration) -> DrainOutcome
+where
+ K: Debug + 'static,
+{
+ use futures::stream::{FuturesUnordered, StreamExt};
+
+ if drained.is_empty() {
+ return DrainOutcome::default();
+ }
+
+ let deadline = Instant::now() + timeout;
+
+ let mut pending: FuturesUnordered<_> = drained
+ .into_iter()
+ .map(|(key, mut entry)| async move {
+ // Closing the sender unblocks the writer task's
+ // `recv().await` with `Err(Closed)` so it exits even if
+ // the shutdown token has not been triggered. Writer is
+ // awaited before the reader handle is dropped so a
+ // mid-writev cancellation cannot truncate a frame.
+ entry.sender.close();
+ // Writer gets the full shared deadline so a legitimately-
+ // flushing `write_vectored_all` never loses budget to any
+ // other entry in the same drain batch.
+ let writer = drain_handle(entry.writer_handle.take(), deadline, &key).await;
+ // Reader gets at least half of what is left after the writer
+ // returns, or READER_DRAIN_FLOOR, whichever is larger. This
+ // prevents a slow writer from force-cancelling the reader
+ // even though reader exit is usually immediate on socket
+ // half-close.
+ let reader_deadline = {
+ let now = Instant::now();
+ let remaining = deadline.saturating_duration_since(now);
+ let reader_budget = std::cmp::max(remaining / 2, READER_DRAIN_FLOOR);
+ now + reader_budget
+ };
+ let reader = drain_handle(entry.reader_handle.take(), reader_deadline, &key).await;
+ (writer, reader)
+ })
+ .collect();
+
+ let mut clean = 0usize;
+ let mut force = 0usize;
+ let mut tally = |outcome: TaskOutcome| match outcome {
+ TaskOutcome::None => {}
+ TaskOutcome::Clean => clean += 1,
+ TaskOutcome::Force => force += 1,
+ };
+ while let Some((writer, reader)) = pending.next().await {
+ tally(writer);
+ tally(reader);
+ }
+
+ DrainOutcome {
+ clean,
+ force,
+ background_clean: 0,
+ background_force: 0,
+ }
+}
+
+#[derive(Debug)]
+enum TaskOutcome {
+ /// No handle to await (already taken).
+ None,
+ /// Task exited within the deadline.
+ Clean,
+ /// Task was force-cancelled at or past the deadline.
+ Force,
+}
+
+#[allow(clippy::future_not_send)] // single-threaded compio
+async fn drain_handle<K: Debug>(
+ handle: Option<JoinHandle<()>>,
+ deadline: Instant,
+ key: &K,
+) -> TaskOutcome {
+ let Some(handle) = handle else {
+ return TaskOutcome::None;
+ };
+ let remaining = deadline.saturating_duration_since(Instant::now());
+ if remaining.is_zero() {
+ debug!(?key, "drain deadline reached, cancelling task");
+ drop(handle);
+ return TaskOutcome::Force;
+ }
+ if compio::time::timeout(remaining, handle).await.is_ok() {
+ TaskOutcome::Clean
+ } else {
+ warn!(?key, "task did not exit within deadline");
+ TaskOutcome::Force
+ }
+}
+
+/// Connection registry specialised for the replica keyspace (`u8`).
+///
+/// The client registry uses `u128` client ids so a `HashMap` backing is
+/// the right trade-off. Replicas are capped at 256 by the keyspace, and
+/// in practice the cluster has 3-7 replicas, so a fixed
+/// `[Option<Entry>; 256]` avoids every send-path lookup paying hash +
+/// bucket indirection. Storage is ~10 KB per bus, paid once at
+/// construction.
+///
+/// The API mirrors [`ConnectionRegistry`] exactly; `IggyMessageBus`
+/// swaps the backing type transparently to every existing call site.
+#[derive(Debug)]
+pub struct ReplicaRegistry {
+ slots: RefCell<[Option<Entry>; 256]>,
+ len: std::cell::Cell<usize>,
+}
+
+impl Default for ReplicaRegistry {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ReplicaRegistry {
+ #[must_use]
+ pub fn new() -> Self {
+ Self {
+ slots: RefCell::new(std::array::from_fn(|_| None)),
+ len: std::cell::Cell::new(0),
+ }
+ }
+
+ #[must_use]
+ pub fn contains(&self, key: u8) -> bool {
+ self.slots.borrow()[usize::from(key)].is_some()
+ }
+
+ #[must_use]
+ pub const fn len(&self) -> usize {
+ self.len.get()
+ }
+
+ #[must_use]
+ pub const fn is_empty(&self) -> bool {
+ self.len.get() == 0
+ }
+
+ /// See [`ConnectionRegistry::borrow_sender`].
+ pub fn borrow_sender(&self, key: u8) -> Option<std::cell::Ref<'_, BusSender>> {
+ std::cell::Ref::filter_map(self.slots.borrow(), |slots| {
+ slots[usize::from(key)].as_ref().map(|entry| &entry.sender)
+ })
+ .ok()
+ }
+
+ /// See [`ConnectionRegistry::insert`].
+ ///
+ /// # Errors
+ ///
+ /// Returns [`AlreadyRegistered`] if `key` is already present.
+ pub fn insert(
+ &self,
+ key: u8,
+ sender: BusSender,
+ writer_handle: JoinHandle<()>,
+ reader_handle: JoinHandle<()>,
+ ) -> Result<(), AlreadyRegistered> {
+ let mut slots = self.slots.borrow_mut();
+ let slot = &mut slots[usize::from(key)];
+ if slot.is_some() {
+ return Err(AlreadyRegistered);
+ }
+ *slot = Some(Entry {
+ sender,
+ writer_handle: Some(writer_handle),
+ reader_handle: Some(reader_handle),
+ });
+ self.len.set(self.len.get() + 1);
+ Ok(())
+ }
+
+ /// See [`ConnectionRegistry::remove`].
+ pub fn remove(&self, key: u8) -> bool {
+ if self.slots.borrow_mut()[usize::from(key)].take().is_some() {
+ self.len.set(self.len.get() - 1);
+ true
+ } else {
+ false
+ }
+ }
+
+ /// See [`ConnectionRegistry::close_peer`].
+ #[allow(clippy::future_not_send)] // single-threaded compio
+ pub async fn close_peer(&self, key: u8, timeout: Duration) {
+ let mut entry = {
+ let mut slots = self.slots.borrow_mut();
+ let Some(entry) = slots[usize::from(key)].take() else {
+ return;
+ };
+ self.len.set(self.len.get() - 1);
+ entry
+ };
+ entry.sender.close();
+ if let Some(writer_handle) = entry.writer_handle.take() {
+ let _ = compio::time::timeout(timeout, writer_handle).await;
+ }
+ drop(entry.reader_handle);
+ }
+
+ /// See [`ConnectionRegistry::drain`].
+ #[allow(clippy::future_not_send)] // single-threaded compio
+ pub async fn drain(&self, timeout: Duration) -> DrainOutcome {
+ let drained: Vec<(u8, Entry)> = {
+ let mut slots = self.slots.borrow_mut();
+ let mut out = Vec::with_capacity(self.len.get());
+ for (idx, slot) in slots.iter_mut().enumerate() {
+ if let Some(entry) = slot.take() {
+ // Safe: index comes from iterating a 256-element array.
+ #[allow(clippy::cast_possible_truncation)]
+ out.push((idx as u8, entry));
+ }
+ }
+ self.len.set(0);
+ out
+ };
+ drain_entries(drained, timeout).await
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::lifecycle::Shutdown;
+ use iggy_binary_protocol::{Command2, GenericHeader, HEADER_SIZE, Message};
+
+ #[allow(clippy::cast_possible_truncation)]
+ fn make_bus_msg() -> BusMessage {
+ Message::<GenericHeader>::new(HEADER_SIZE)
+ .transmute_header(|_, h: &mut GenericHeader| {
+ h.command = Command2::Ping;
+ h.size = HEADER_SIZE as u32;
+ })
+ .into_frozen()
+ }
+
+ fn spawn_dummy_writer(rx: BusReceiver) -> JoinHandle<()> {
+ compio::runtime::spawn(async move {
+ while let Ok(_msg) = rx.recv().await {
+ // discard
+ }
+ })
+ }
+
+ fn spawn_dummy_reader(token: crate::lifecycle::ShutdownToken) -> JoinHandle<()> {
+ compio::runtime::spawn(async move {
+ token.wait().await;
+ })
+ }
+
+ #[compio::test]
+ async fn insert_and_get_sender() {
+ let reg: ConnectionRegistry<u8> = ConnectionRegistry::new();
+ let (_shutdown, token) = Shutdown::new();
+ let (tx, rx) = async_channel::bounded(8);
+ let writer = spawn_dummy_writer(rx);
+ let reader = spawn_dummy_reader(token);
+
+ reg.insert(1u8, tx, writer, reader).expect("insert ok");
+ assert!(reg.contains(1u8));
+ assert_eq!(reg.len(), 1);
+
+ let sender = reg.borrow_sender(1u8).expect("sender present");
+ sender.try_send(make_bus_msg()).expect("queue accepts msg");
+ }
+
+ #[compio::test]
+ async fn insert_duplicate_errors() {
+ let reg: ConnectionRegistry<u8> = ConnectionRegistry::new();
+ let (_shutdown, token) = Shutdown::new();
+ let (tx1, rx1) = async_channel::bounded(8);
+ let (tx2, rx2) = async_channel::bounded(8);
+ let w1 = spawn_dummy_writer(rx1);
+ let r1 = spawn_dummy_reader(token.clone());
+ let w2 = spawn_dummy_writer(rx2);
+ let r2 = spawn_dummy_reader(token);
+
+ reg.insert(1u8, tx1, w1, r1).expect("first insert");
+ let err = reg.insert(1u8, tx2, w2, r2);
+ assert!(err.is_err());
+ }
+
+ #[compio::test]
+ async fn drain_after_shutdown_counts_both_tasks_per_entry() {
+ let reg: ConnectionRegistry<u8> = ConnectionRegistry::new();
+ let (shutdown, token) = Shutdown::new();
+
+ for k in 0..3u8 {
+ let (tx, rx) = async_channel::bounded(8);
+ let w = spawn_dummy_writer(rx);
+ let r = spawn_dummy_reader(token.clone());
+ reg.insert(k, tx, w, r).unwrap();
+ }
+
+ shutdown.trigger();
+ let outcome = reg.drain(Duration::from_secs(2)).await;
+ // 3 entries * 2 tasks (writer + reader) each = 6 clean exits.
+ assert_eq!(outcome.clean, 6);
+ assert_eq!(outcome.force, 0);
+ assert!(reg.is_empty());
+ }
+
+ #[compio::test]
+ async fn drain_force_cancels_reader_that_refuses_to_exit() {
+ let reg: ConnectionRegistry<u8> = ConnectionRegistry::new();
+ let (_shutdown, _token) = Shutdown::new();
+ let (tx, rx) = async_channel::bounded(8);
+ let w = spawn_dummy_writer(rx);
+ // Reader ignores shutdown entirely.
+ let r = compio::runtime::spawn(async move {
+ loop {
+ compio::time::sleep(Duration::from_secs(10)).await;
+ }
+ });
+ reg.insert(1u8, tx, w, r).unwrap();
+
+ let outcome = reg.drain(Duration::from_millis(40)).await;
+ // Writer exits because the sender is closed by drain. Reader is
+ // force-cancelled.
+ assert_eq!(outcome.clean, 1);
+ assert_eq!(outcome.force, 1);
+ }
+
+ #[compio::test]
+ async fn close_peer_closes_sender_awaits_writer() {
+ let reg: ConnectionRegistry<u8> = ConnectionRegistry::new();
+ let (_shutdown, token) = Shutdown::new();
+ let (tx, rx) = async_channel::bounded(8);
+ let writer = spawn_dummy_writer(rx);
+ let reader = spawn_dummy_reader(token);
+ reg.insert(7u8, tx, writer, reader).expect("insert ok");
+
+ assert!(reg.contains(7u8));
+ reg.close_peer(7u8, Duration::from_secs(1)).await;
+ assert!(!reg.contains(7u8));
+ }
+
+ #[compio::test]
+ async fn close_peer_noop_on_missing_key() {
+ let reg: ConnectionRegistry<u8> = ConnectionRegistry::new();
+ reg.close_peer(42u8, Duration::from_millis(10)).await;
+ }
+
+ /// Writer consumes nearly the whole shared deadline; reader still
+ /// gets at least `READER_DRAIN_FLOOR` of runway and exits cleanly.
+ /// Before the split the reader would be force-cancelled with zero
+ /// remaining time.
+ #[compio::test]
+ async fn drain_reader_keeps_floor_when_writer_slow() {
+ let reg: ConnectionRegistry<u8> = ConnectionRegistry::new();
+ let (shutdown, token) = Shutdown::new();
+
+ let (tx, rx) = async_channel::bounded::<BusMessage>(8);
+ // Writer sleeps just under the shared deadline before returning.
+ let writer = compio::runtime::spawn(async move {
+ while rx.recv().await.is_ok() {}
+ compio::time::sleep(Duration::from_millis(180)).await;
+ });
+ let reader = spawn_dummy_reader(token);
+ reg.insert(1u8, tx, writer, reader).unwrap();
+
+ shutdown.trigger();
+ // Shared deadline = 200 ms. Writer consumes ~180 ms, leaving
+ // ~20 ms. Without the reader floor that would force the reader.
+ let outcome = reg.drain(Duration::from_millis(200)).await;
+ assert_eq!(outcome.force, 0, "reader should drain cleanly: {outcome:?}");
+ assert_eq!(outcome.clean, 2);
+ }
+
+ /// Every entry's writer takes ~80 ms to exit (simulating a slow
+ /// in-flight batch). With sequential drain that would be N * 80 ms.
+ /// With parallel drain it should be bounded by the slowest entry.
+ #[compio::test]
+ async fn drain_runs_entries_concurrently() {
+ const N: u8 = 10;
+ const PER_ENTRY_LATENCY: Duration = Duration::from_millis(80);
+
+ let reg: ConnectionRegistry<u8> = ConnectionRegistry::new();
+ let (shutdown, token) = Shutdown::new();
+
+ for k in 0..N {
+ let (tx, rx) = async_channel::bounded::<BusMessage>(8);
+ // Writer waits PER_ENTRY_LATENCY after observing channel
+ // close before returning, mimicking a slow tail batch.
+ let writer = compio::runtime::spawn(async move {
+ while rx.recv().await.is_ok() {}
+ compio::time::sleep(PER_ENTRY_LATENCY).await;
+ });
+ let reader = spawn_dummy_reader(token.clone());
+ reg.insert(k, tx, writer, reader).unwrap();
+ }
+
+ // Shutdown so readers (which wait on the token) exit cleanly.
+ shutdown.trigger();
+
+ let start = Instant::now();
+ let outcome = reg.drain(Duration::from_secs(5)).await;
+ let elapsed = start.elapsed();
+
+ assert_eq!(outcome.clean, usize::from(N) * 2);
+ assert_eq!(outcome.force, 0);
+ // Sequential lower bound would be N * PER_ENTRY_LATENCY.
+ // Allow 3x the single-entry latency as headroom for scheduling.
+ let parallel_budget = PER_ENTRY_LATENCY * 3;
+ assert!(
+ elapsed < parallel_budget,
+ "drain took {:?}, expected parallel < {:?} (serial would be ~{:?})",
+ elapsed,
+ parallel_budget,
+ PER_ENTRY_LATENCY * u32::from(N),
+ );
+ }
+}
diff --git a/core/message_bus/src/cache/mod.rs b/core/message_bus/src/lifecycle/mod.rs
similarity index 63%
copy from core/message_bus/src/cache/mod.rs
copy to core/message_bus/src/lifecycle/mod.rs
index fd239df97..2ceae3107 100644
--- a/core/message_bus/src/cache/mod.rs
+++ b/core/message_bus/src/lifecycle/mod.rs
@@ -15,16 +15,17 @@
// specific language governing permissions and limitations
// under the License.
-use crate::cache::connection::ShardedState;
+//! Bus lifecycle primitives.
+//!
+//! A root cancellation token and a keyed connection registry. Intentionally
+//! narrow: these are the only two abstractions we need to drive graceful
+//! shutdown for the TCP path (IGGY-99).
-// TODO: Move to some common trait location.
-/// Allocation strategy that produces deltas for a specific sharded state type
-pub trait AllocationStrategy<SS>
-where
- SS: ShardedState,
-{
- fn allocate(&self, entry: SS::Entry) -> Option<SS::Delta>;
- fn deallocate(&self, entry: SS::Entry) -> Option<SS::Delta>;
-}
+pub mod connection_registry;
+pub mod shutdown;
-pub mod connection;
+pub use connection_registry::{
+ AlreadyRegistered, BusMessage, BusReceiver, BusSender, ConnectionRegistry, DrainOutcome,
+ ReplicaRegistry,
+};
+pub use shutdown::{Shutdown, ShutdownToken};
diff --git a/core/message_bus/src/lifecycle/shutdown.rs b/core/message_bus/src/lifecycle/shutdown.rs
new file mode 100644
index 000000000..604836191
--- /dev/null
+++ b/core/message_bus/src/lifecycle/shutdown.rs
@@ -0,0 +1,184 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Root shutdown token for the message bus.
+//!
+//! A single [`Shutdown`] is paired with any number of cloned [`ShutdownToken`]s.
+//! Tasks (accept loops, read loops, reconnect periodics) keep a clone and
+//! `select!` on [`ShutdownToken::wait`] or poll [`ShutdownToken::is_triggered`].
+//! When [`Shutdown::trigger`] is called, every clone observes the transition
+//! in the same runtime turn.
+//!
+//! Single-threaded compio means `Rc<Cell<bool>>` is correct; no `Arc`/atomics.
+
+use async_channel::{Receiver, Sender};
+use futures::FutureExt;
+use std::cell::Cell;
+use std::rc::Rc;
+use std::time::Duration;
+
+/// Owning side of the shutdown signal.
+///
+/// Create via [`Shutdown::new`], hand out [`ShutdownToken`]s via [`ShutdownToken::clone`],
+/// and call [`Shutdown::trigger`] exactly once to stop everyone.
+#[derive(Debug)]
+pub struct Shutdown {
+ sender: Sender<()>,
+ triggered: Rc<Cell<bool>>,
+}
+
+/// Cheaply cloneable observer for a [`Shutdown`].
+#[derive(Debug, Clone)]
+pub struct ShutdownToken {
+ receiver: Receiver<()>,
+ triggered: Rc<Cell<bool>>,
+}
+
+impl Shutdown {
+ /// Create a new shutdown signal and its first observer token.
+ #[must_use]
+ pub fn new() -> (Self, ShutdownToken) {
+ let (sender, receiver) = async_channel::bounded::<()>(1);
+ let triggered = Rc::new(Cell::new(false));
+ let shutdown = Self {
+ sender,
+ triggered: triggered.clone(),
+ };
+ let token = ShutdownToken {
+ receiver,
+ triggered,
+ };
+ (shutdown, token)
+ }
+
+ /// Flip the bit and close the broadcast channel.
+ ///
+ /// Idempotent: a second call is a no-op. Does not consume `self` so
+ /// the bus can keep the `Shutdown` as a plain field and trigger via
+ /// a shared reference.
+ pub fn trigger(&self) {
+ if !self.triggered.replace(true) {
+ self.sender.close();
+ }
+ }
+
+ /// Whether [`trigger`](Self::trigger) has been called.
+ #[must_use]
+ pub fn is_triggered(&self) -> bool {
+ self.triggered.get()
+ }
+}
+
+impl ShutdownToken {
+ /// O(1) non-blocking check. True after [`Shutdown::trigger`] returns.
+ #[must_use]
+ pub fn is_triggered(&self) -> bool {
+ self.triggered.get()
+ }
+
+ /// Resolves the instant shutdown fires.
+ ///
+ /// Safe to poll concurrently from many tasks because `async_channel`
+ /// receivers resolve to `Err(Closed)` for every waiter once the sender
+ /// side closes.
+ #[allow(clippy::future_not_send)]
+ pub async fn wait(&self) {
+ // `recv()` on a closed channel returns immediately with an error,
+ // so this loop exits in one await as soon as the sender side closes.
+ while self.receiver.recv().await.is_ok() {}
+ }
+
+ /// Sleep for `duration`, returning early if shutdown fires.
+ ///
+ /// Returns `true` if the full duration elapsed, `false` if shutdown
+ /// interrupted the sleep. Used by periodic tasks:
+ ///
+ /// ```ignore
+ /// while token.sleep_or_shutdown(PERIOD).await {
+ /// tick().await;
+ /// }
+ /// ```
+ #[allow(clippy::future_not_send)]
+ pub async fn sleep_or_shutdown(&self, duration: Duration) -> bool {
+ if self.is_triggered() {
+ return false;
+ }
+ futures::select! {
+ () = self.wait().fuse() => false,
+ () = compio::time::sleep(duration).fuse() => !self.is_triggered(),
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use std::time::Duration;
+
+ #[compio::test]
+ async fn trigger_sets_flag_and_resolves_wait() {
+ let (shutdown, token) = Shutdown::new();
+ assert!(!token.is_triggered());
+
+ shutdown.trigger();
+ assert!(token.is_triggered());
+ token.wait().await; // resolves immediately
+ }
+
+ #[compio::test]
+ async fn wait_resolves_for_all_clones() {
+ let (shutdown, token_a) = Shutdown::new();
+ let token_b = token_a.clone();
+
+ shutdown.trigger();
+ token_a.wait().await;
+ token_b.wait().await;
+ assert!(token_a.is_triggered());
+ assert!(token_b.is_triggered());
+ }
+
+ #[compio::test]
+ async fn sleep_or_shutdown_returns_true_on_elapsed() {
+ let (_shutdown, token) = Shutdown::new();
+ let ok = token.sleep_or_shutdown(Duration::from_millis(20)).await;
+ assert!(ok);
+ assert!(!token.is_triggered());
+ }
+
+ #[compio::test]
+ async fn sleep_or_shutdown_returns_false_when_triggered_first() {
+ let (shutdown, token) = Shutdown::new();
+ // Already triggered
+ shutdown.trigger();
+ let ok = token.sleep_or_shutdown(Duration::from_secs(10)).await;
+ assert!(!ok);
+ }
+
+ #[compio::test]
+ async fn sleep_or_shutdown_interrupted_mid_sleep() {
+ let (shutdown, token) = Shutdown::new();
+ let token_clone = token.clone();
+ let sleeper = compio::runtime::spawn(async move {
+ token_clone.sleep_or_shutdown(Duration::from_secs(10)).await
+ });
+ // Let the sleeper reach the select
+ compio::time::sleep(Duration::from_millis(5)).await;
+ shutdown.trigger();
+ let ok = sleeper.await.expect("task ok");
+ assert!(!ok);
+ }
+}
diff --git a/core/message_bus/src/replica_io.rs b/core/message_bus/src/replica_io.rs
new file mode 100644
index 000000000..5ef866ff3
--- /dev/null
+++ b/core/message_bus/src/replica_io.rs
@@ -0,0 +1,154 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Shard-0 TCP plane bootstrap.
+//!
+//! Every shard on a node instantiates its own `Rc<IggyMessageBus>` on its
+//! own compio runtime, but only shard 0 terminates TCP: it binds the
+//! replica listener, binds the client listener, and dials higher-id peers.
+//! Each accepted or dialed connection is handed to the delegate callbacks
+//! supplied by the caller (typically wrapping `shard::coordinator::ShardZeroCoordinator`),
+//! which duplicate the fd and ship it to the owning shard via the
+//! inter-shard `ShardFrame` channel.
+//!
+//! 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.
+
+use std::net::SocketAddr;
+use std::rc::Rc;
+use std::time::Duration;
+
+use iggy_common::IggyError;
+
+use crate::connector::start as start_connector;
+use crate::replica_listener::{bind as bind_replica_listener, run as run_replica_listener};
+use crate::{AcceptedClientFn, AcceptedReplicaFn, IggyMessageBus, client_listener};
+
+/// Bound addresses returned to shard 0 after the listeners come up.
+#[derive(Debug, Clone, Copy)]
+pub struct BoundPlanes {
+ pub replica: SocketAddr,
+ pub client: SocketAddr,
+}
+
+/// Bind the replica + client listeners on shard 0 and start the outbound
+/// replica connector. Non-zero shards early-return `Ok(None)`.
+///
+/// Each accepted / dialed TCP stream is handed to the supplied delegate
+/// callback: `on_accepted_replica` for inbound replicas and for freshly
+/// dialed higher-id peers; `on_accepted_client` for SDK client accepts.
+/// The callbacks are responsible for the dup-fd + inter-shard send.
+///
+/// # Errors
+///
+/// Returns `IggyError::CannotBindToSocket` if either listener bind fails.
+#[allow(clippy::future_not_send)]
+#[allow(clippy::too_many_arguments)]
+pub async fn start_on_shard_zero(
+ bus: &Rc<IggyMessageBus>,
+ replica_listen_addr: SocketAddr,
+ client_listen_addr: SocketAddr,
+ cluster_id: u128,
+ self_id: u8,
+ replica_count: u8,
+ peers: Vec<(u8, SocketAddr)>,
+ on_accepted_replica: AcceptedReplicaFn,
+ on_accepted_client: AcceptedClientFn,
+ reconnect_period: Duration,
+) -> Result<Option<BoundPlanes>, IggyError> {
+ if bus.shard_id() != 0 {
+ return Ok(None);
+ }
+
+ let (replica_listener, replica_bound) = bind_replica_listener(replica_listen_addr).await?;
+ let (clients_listener, client_bound) = client_listener::bind(client_listen_addr).await?;
+
+ let token_for_replica = bus.token();
+ let on_accepted_replica_for_listener = on_accepted_replica.clone();
+ let listener_max_message_size = bus.config().max_message_size;
+ let replica_handle = compio::runtime::spawn(async move {
+ run_replica_listener(
+ replica_listener,
+ token_for_replica,
+ cluster_id,
+ self_id,
+ replica_count,
+ on_accepted_replica_for_listener,
+ listener_max_message_size,
+ )
+ .await;
+ });
+ bus.track_background(replica_handle);
+
+ let token_for_client = bus.token();
+ let client_handle = compio::runtime::spawn(async move {
+ client_listener::run(clients_listener, token_for_client, on_accepted_client).await;
+ });
+ bus.track_background(client_handle);
+
+ start_connector(
+ bus,
+ cluster_id,
+ self_id,
+ peers,
+ on_accepted_replica,
+ reconnect_period,
+ )
+ .await;
+
+ Ok(Some(BoundPlanes {
+ replica: replica_bound,
+ client: client_bound,
+ }))
+}
+
+/// [`start_on_shard_zero`] defaulting `reconnect_period` to the bus's
+/// [`crate::MessageBusConfig::reconnect_period`].
+///
+/// # Errors
+///
+/// Returns `IggyError::CannotBindToSocket` if either listener bind fails.
+#[allow(clippy::future_not_send)]
+#[allow(clippy::too_many_arguments)]
+pub async fn start_on_shard_zero_default(
+ bus: &Rc<IggyMessageBus>,
+ replica_listen_addr: SocketAddr,
+ client_listen_addr: SocketAddr,
+ cluster_id: u128,
+ self_id: u8,
+ replica_count: u8,
+ peers: Vec<(u8, SocketAddr)>,
+ on_accepted_replica: AcceptedReplicaFn,
+ on_accepted_client: AcceptedClientFn,
+) -> Result<Option<BoundPlanes>, IggyError> {
+ let reconnect_period = bus.config().reconnect_period;
+ start_on_shard_zero(
+ bus,
+ replica_listen_addr,
+ client_listen_addr,
+ cluster_id,
+ self_id,
+ replica_count,
+ peers,
+ on_accepted_replica,
+ on_accepted_client,
+ reconnect_period,
+ )
+ .await
+}
diff --git a/core/message_bus/src/replica_listener.rs b/core/message_bus/src/replica_listener.rs
new file mode 100644
index 000000000..9d9db9192
--- /dev/null
+++ b/core/message_bus/src/replica_listener.rs
@@ -0,0 +1,164 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Inbound TCP listener for replica-to-replica consensus traffic.
+//!
+//! Runs only on shard 0. On every successful `Ping` handshake the listener
+//! hands the accepted `TcpStream` to an `on_accepted` callback provided by
+//! the shard bootstrap, which dup-and-ships the fd to the owning shard via
+//! the inter-shard channel (see `shard::coordinator::ShardZeroCoordinator`).
+//! This module no longer installs writer / reader tasks itself.
+//!
+//! Duplicate connections are eliminated by directionality: each replica
+//! only dials peers with strictly greater ids and only accepts inbound
+//! from peers with strictly lower ids. No race, no tiebreaker.
+//!
+//! # Security
+//!
+//! The `Ping` handshake validates `cluster_id`, the directional bound, and
+//! `replica_count`. There is NO shared secret, no mTLS, no version
+//! negotiation. Deploy the replica listener on a trusted network boundary
+//! (cluster-local VPC, private subnet, overlay, or equivalent) - NEVER
+//! expose the replica port directly to the public internet. Follow-up work
+//! on an authenticated handshake is tracked under IGGY-112.
+
+use crate::framing;
+use crate::lifecycle::ShutdownToken;
+use crate::{AcceptedReplicaFn, GenericHeader, Message};
+use compio::net::{SocketOpts, TcpListener, TcpStream};
+use futures::FutureExt;
+use iggy_binary_protocol::Command2;
+use iggy_common::IggyError;
+use std::net::SocketAddr;
+use std::rc::Rc;
+use tracing::{debug, error, info, warn};
+
+/// Handler for inbound replica consensus messages.
+///
+/// Preserved for callers (tests, simulator-facing glue) that want to install
+/// a connection locally without going through the coordinator. The shard-0
+/// production path uses [`AcceptedReplicaFn`] instead.
+pub type MessageHandler = Rc<dyn Fn(u8, Message<GenericHeader>)>;
+
+/// Bind the replica listener and return the bound address.
+///
+/// # Errors
+///
+/// Returns [`IggyError::CannotBindToSocket`] if the bind fails.
+#[allow(clippy::future_not_send)]
+pub async fn bind(addr: SocketAddr) -> Result<(TcpListener, SocketAddr), IggyError> {
+ // `SO_REUSEPORT` intentionally not set: only shard 0 binds the replica
+ // listener. Kernel-level accept distribution would fight the shard-0
+ // coordinator's explicit round-robin allocation.
+ let opts = SocketOpts::new().nodelay(true).keepalive(true);
+ let listener = TcpListener::bind_with_options(addr, &opts)
+ .await
+ .map_err(|_| IggyError::CannotBindToSocket(addr.to_string()))?;
+ let actual = listener
+ .local_addr()
+ .map_err(|e| IggyError::IoError(e.to_string()))?;
+ Ok((listener, actual))
+}
+
+/// Run the inbound replica listener accept loop until the shutdown token
+/// fires. Every successful handshake fires the `on_accepted` callback; the
+/// callback owns the accepted stream from that point on.
+#[allow(clippy::future_not_send)]
+#[allow(clippy::too_many_arguments)]
+pub async fn run(
+ listener: TcpListener,
+ token: ShutdownToken,
+ cluster_id: u128,
+ self_id: u8,
+ replica_count: u8,
+ on_accepted: AcceptedReplicaFn,
+ max_message_size: usize,
+) {
+ info!(
+ "Replica listener accepting on {:?}",
+ listener.local_addr().ok()
+ );
+ loop {
+ futures::select! {
+ () = token.wait().fuse() => {
+ debug!("Replica listener shutting down");
+ break;
+ }
+ result = listener.accept().fuse() => {
+ match result {
+ Ok((mut stream, peer_addr)) => {
+ match handshake(
+ &mut stream,
+ cluster_id,
+ self_id,
+ replica_count,
+ max_message_size,
+ )
+ .await {
+ Ok(peer_id) => {
+ on_accepted(stream, peer_id);
+ }
+ Err(e) => {
+ warn!(%peer_addr, "replica handshake failed: {e}");
+ }
+ }
+ }
+ Err(e) => {
+ error!("Replica listener accept failed: {e}");
+ }
+ }
+ }
+ }
+ }
+}
+
+/// Read and validate the `Ping` handshake from an inbound peer connection.
+///
+/// Returns the peer replica id on success.
+//
+// TODO(IGGY-112): authenticate the replica handshake.
+//
+// Current validation covers `cluster_id` and the directional id bound
+// only: no shared secret, no mTLS, no protocol version negotiation.
+// Any peer that reaches the replica port and knows the `cluster_id`
+// can pose as a replica. Acceptable on the assumed trusted network
+// boundary; hardening is tracked under IGGY-112 (authenticated
+// handshake + version negotiation).
+#[allow(clippy::future_not_send)]
+async fn handshake(
+ stream: &mut TcpStream,
+ our_cluster: u128,
+ self_id: u8,
+ replica_count: u8,
+ max_message_size: usize,
+) -> Result<u8, IggyError> {
+ let msg = framing::read_message(stream, max_message_size).await?;
+ let header = msg.header();
+ if header.command != Command2::Ping {
+ return Err(IggyError::InvalidCommand);
+ }
+ if header.cluster != our_cluster {
+ return Err(IggyError::InvalidCommand);
+ }
+ // Directional rule: a replica only accepts inbound from peers with
+ // strictly lower ids. The peer is responsible for not dialing us if
+ // it has the higher id; this is just defensive.
+ if header.replica >= replica_count || header.replica >= self_id {
+ return Err(IggyError::InvalidCommand);
+ }
+ Ok(header.replica)
+}
diff --git a/core/message_bus/src/socket_opts.rs b/core/message_bus/src/socket_opts.rs
new file mode 100644
index 000000000..e550e6902
--- /dev/null
+++ b/core/message_bus/src/socket_opts.rs
@@ -0,0 +1,59 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Per-connection socket options.
+//!
+//! `compio::net::SocketOpts` exposes `keepalive(bool)` only - it does not
+//! let callers tune `TCP_KEEPIDLE`, `TCP_KEEPINTVL`, or `TCP_KEEPCNT`. We
+//! reach through to `socket2::SockRef` for the finer knobs so the bus can
+//! detect a silently dead peer within a predictable bound without waiting
+//! for the kernel defaults (2 hours idle on Linux).
+
+use compio::net::TcpStream;
+use socket2::{SockRef, TcpKeepalive};
+use std::io;
+use std::time::Duration;
+
+/// Configure TCP keepalive on a freshly accepted or freshly dialed
+/// connection. Uses `socket2` directly because compio's `SocketOpts` only
+/// exposes the on/off toggle; we need the idle/interval/count tuning to get
+/// a bounded detection window on a half-dead peer.
+///
+/// `idle` maps to `TCP_KEEPIDLE`, `interval` to `TCP_KEEPINTVL`, `retries`
+/// to `TCP_KEEPCNT`. Total detection window is roughly
+/// `idle + retries * interval`; keep it comfortably shorter than the VSR
+/// view-change timeout so replicas can observe a dead peer ahead of the
+/// consensus-level decision.
+///
+/// # Errors
+///
+/// Returns the underlying `io::Error` if the kernel rejects any of the
+/// setsockopt calls. Callers log-and-continue rather than tear down the
+/// connection, because a missing keepalive is a soft failure.
+pub fn apply_keepalive_for_connection(
+ stream: &TcpStream,
+ idle: Duration,
+ interval: Duration,
+ retries: u32,
+) -> io::Result<()> {
+ let sock = SockRef::from(stream);
+ let params = TcpKeepalive::new()
+ .with_time(idle)
+ .with_interval(interval)
+ .with_retries(retries);
+ sock.set_tcp_keepalive(¶ms)
+}
diff --git a/core/message_bus/src/cache/mod.rs b/core/message_bus/src/transports/mod.rs
similarity index 67%
copy from core/message_bus/src/cache/mod.rs
copy to core/message_bus/src/transports/mod.rs
index fd239df97..aa0c7c27b 100644
--- a/core/message_bus/src/cache/mod.rs
+++ b/core/message_bus/src/transports/mod.rs
@@ -15,16 +15,8 @@
// specific language governing permissions and limitations
// under the License.
-use crate::cache::connection::ShardedState;
-
-// TODO: Move to some common trait location.
-/// Allocation strategy that produces deltas for a specific sharded state type
-pub trait AllocationStrategy<SS>
-where
- SS: ShardedState,
-{
- fn allocate(&self, entry: SS::Entry) -> Option<SS::Delta>;
- fn deallocate(&self, entry: SS::Entry) -> Option<SS::Delta>;
-}
-
-pub mod connection;
+//! Transport abstraction - placeholder.
+//!
+//! Tracked by IGGY-112. This module is deliberately empty until the
+//! `Transport` / `TransportListener` / `TransportConn` trait family lands;
+//! do not grow speculative scaffolding here.
diff --git a/core/message_bus/src/writer_task.rs b/core/message_bus/src/writer_task.rs
new file mode 100644
index 000000000..d1219ae47
--- /dev/null
+++ b/core/message_bus/src/writer_task.rs
@@ -0,0 +1,100 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Per-connection writer task.
+//!
+//! Drains the per-peer mpsc and writes batched consensus frames in a single
+//! `writev` syscall via `compio::io::AsyncWriteExt::write_vectored_all`.
+//!
+//! Single-producer (the bus `send_to_*` path) -> single-consumer (this task),
+//! per peer. The mpsc bound provides backpressure: callers `try_send` and
+//! receive [`SendError::Backpressure`](crate::SendError::Backpressure) when
+//! the queue is full. VSR retransmits dropped messages from the WAL.
+//!
+//! Lifecycle: the task exits cleanly when either
+//! - the bus shutdown token fires (preferred path during graceful shutdown),
+//! - the `Sender` side is dropped/closed (peer unregistered), or
+//! - a write to the wire fails (broken connection).
+
+use crate::lifecycle::{BusMessage, BusReceiver, ShutdownToken};
+use compio::io::AsyncWriteExt;
+use compio::net::{OwnedWriteHalf, TcpStream};
+use futures::FutureExt;
+use tracing::{debug, trace, warn};
+
+/// Run the per-connection writer loop until the channel closes, the
+/// shutdown token fires, or a write fails.
+///
+/// `max_batch` caps how many messages a single `writev` syscall coalesces.
+/// Larger batches reduce syscalls per N messages at the cost of memory
+/// per batch and worst-case latency for the head-of-batch message.
+#[allow(clippy::future_not_send)] // single-threaded compio
+pub async fn run(
+ rx: BusReceiver,
+ mut write_half: OwnedWriteHalf<TcpStream>,
+ token: ShutdownToken,
+ label: &'static str,
+ peer: String,
+ max_batch: usize,
+) {
+ let mut batch: Vec<BusMessage> = Vec::with_capacity(max_batch);
+
+ loop {
+ let first = futures::select! {
+ () = token.wait().fuse() => {
+ debug!(%label, %peer, "writer task: shutdown token fired");
+ return;
+ }
+ msg = rx.recv().fuse() => {
+ if let Ok(m) = msg {
+ m
+ } else {
+ debug!(%label, %peer, "writer task: channel closed");
+ return;
+ }
+ }
+ };
+
+ batch.push(first);
+
+ // Drain the rest non-blocking up to max_batch so a single writev
+ // syscall covers as many messages as the queue currently holds.
+ while batch.len() < max_batch {
+ match rx.try_recv() {
+ Ok(m) => batch.push(m),
+ Err(_) => break,
+ }
+ }
+
+ let drained = batch.len();
+ trace!(%label, %peer, batch = drained, "writev batch");
+
+ // Single writev for the whole batch. write_vectored_all loops
+ // internally on partial writes until the full batch lands or the
+ // socket errors.
+ let to_write = std::mem::take(&mut batch);
+ let compio::BufResult(result, mut returned) = write_half.write_vectored_all(to_write).await;
+ if let Err(e) = result {
+ warn!(%label, %peer, error = ?e, "writer task: write_vectored_all failed");
+ return;
+ }
+
+ // Reuse the Vec allocation across iterations.
+ returned.clear();
+ batch = returned;
+ }
+}
diff --git a/core/message_bus/tests/backpressure.rs b/core/message_bus/tests/backpressure.rs
new file mode 100644
index 000000000..1a399aa04
--- /dev/null
+++ b/core/message_bus/tests/backpressure.rs
@@ -0,0 +1,114 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Verify that the per-peer queue applies backpressure: with a tiny capacity
+//! and a peer whose receiving end never drains, `send_to_replica` returns
+//! `SendError::Backpressure` once the queue is full.
+//!
+//! This proves the bus is fire-and-forget and that VSR can rely on
+//! `Backpressure` as a drop signal rather than a blocking await.
+
+mod common;
+
+use common::{header_only, install_replicas_locally, loopback};
+use iggy_binary_protocol::Command2;
+use message_bus::connector::{DEFAULT_RECONNECT_PERIOD, start as start_connector};
+use message_bus::replica_listener::{MessageHandler, bind, run};
+use message_bus::{IggyMessageBus, MessageBus, SendError};
+use std::rc::Rc;
+use std::time::Duration;
+
+const CLUSTER: u128 = 0xDEAD;
+
+#[compio::test]
+async fn try_send_returns_backpressure_when_queue_full() {
+ // Tiny per-peer queue so it fills after a few sends.
+ let bus0 = Rc::new(IggyMessageBus::with_capacity(0, 4));
+
+ // Peer 1: a real listener that NEVER reads its socket. The kernel TCP
+ // recv buffer fills, then the writer task on bus0 blocks on its next
+ // writev, which leaves the per-peer mpsc filling up until try_send
+ // returns Full.
+ let bus1 = Rc::new(IggyMessageBus::with_capacity(0, 4));
+ let on_message1: MessageHandler = Rc::new(|_, _| {
+ // Block the read loop forever by sleeping. The handler is sync,
+ // so no actual blocking is possible from here; instead we just do
+ // nothing and rely on the kernel TCP buffer + bus0's tiny queue
+ // to apply backpressure.
+ });
+ let (l1, addr1) = bind(loopback()).await.unwrap();
+ let token_for_l1 = bus1.token();
+ let accept_1 = install_replicas_locally(bus1.clone(), on_message1);
+ let l1_handle = compio::runtime::spawn(async move {
+ run(
+ l1,
+ token_for_l1,
+ CLUSTER,
+ 1,
+ 2,
+ accept_1,
+ message_bus::framing::MAX_MESSAGE_SIZE,
+ )
+ .await;
+ });
+ bus1.track_background(l1_handle);
+
+ // bus0 dials bus1.
+ let on_message0: MessageHandler = Rc::new(|_, _| {});
+ let dial_0 = install_replicas_locally(bus0.clone(), on_message0);
+ start_connector(
+ &bus0,
+ CLUSTER,
+ 0,
+ vec![(1, addr1)],
+ dial_0,
+ DEFAULT_RECONNECT_PERIOD,
+ )
+ .await;
+
+ // Wait for the connection to register.
+ let deadline = std::time::Instant::now() + Duration::from_secs(2);
+ while !bus0.replicas().contains(1) {
+ assert!(std::time::Instant::now() < deadline, "connect timeout");
+ compio::time::sleep(Duration::from_millis(5)).await;
+ }
+
+ // Hammer the bus with messages much bigger than the kernel TCP send
+ // buffer plus the per-peer queue. With queue capacity 4 and the writer
+ // task absorbing some before getting blocked on the wire, we should
+ // see Backpressure within a few thousand sends. We are not measuring
+ // the exact threshold, only that it is finite and reachable.
+ let mut hit_backpressure = false;
+ for _ in 0..100_000 {
+ let msg = header_only(Command2::Prepare, CLUSTER, 0);
+ match bus0.send_to_replica(1, msg.into_frozen()).await {
+ Ok(()) => {}
+ Err(SendError::Backpressure) => {
+ hit_backpressure = true;
+ break;
+ }
+ Err(other) => panic!("unexpected error: {other}"),
+ }
+ }
+ assert!(
+ hit_backpressure,
+ "expected SendError::Backpressure under sustained send load"
+ );
+
+ bus0.shutdown(Duration::from_secs(2)).await;
+ bus1.shutdown(Duration::from_secs(2)).await;
+}
diff --git a/core/message_bus/tests/common/mod.rs b/core/message_bus/tests/common/mod.rs
new file mode 100644
index 000000000..b00fbb310
--- /dev/null
+++ b/core/message_bus/tests/common/mod.rs
@@ -0,0 +1,88 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Shared helpers for integration tests of the `message_bus` crate.
+//!
+//! Each integration test binary in `tests/` includes this module via
+//! `mod common;` and uses the helpers to keep the tests compact.
+//!
+//! Note: integration tests live in separate binaries, so each binary that
+//! does not consume every helper triggers `dead_code`. The crate-level
+//! `#![allow(dead_code)]` below is intentional.
+
+#![allow(dead_code)] // each test binary uses a subset
+
+use iggy_binary_protocol::{Command2, GenericHeader, HEADER_SIZE, Message};
+use message_bus::client_listener::RequestHandler;
+use message_bus::replica_listener::MessageHandler;
+use message_bus::{AcceptedClientFn, AcceptedReplicaFn, IggyMessageBus, installer};
+use std::cell::Cell;
+use std::net::SocketAddr;
+use std::rc::Rc;
+
+/// Loopback address with OS-chosen port.
+#[must_use]
+pub fn loopback() -> SocketAddr {
+ "127.0.0.1:0".parse().unwrap()
+}
+
+/// Build a header-only consensus message with the given command.
+///
+/// Used by tests to fabricate `Request`, `Reply`, `Ping`, etc. directly.
+#[must_use]
+#[allow(clippy::cast_possible_truncation)]
+pub fn header_only(command: Command2, cluster: u128, replica: u8) -> Message<GenericHeader> {
+ Message::<GenericHeader>::new(HEADER_SIZE).transmute_header(|_, h: &mut GenericHeader| {
+ h.command = command;
+ h.cluster = cluster;
+ h.replica = replica;
+ h.size = HEADER_SIZE as u32;
+ })
+}
+
+/// Build an [`AcceptedReplicaFn`] that installs accepted replica streams
+/// directly on the given bus. Mimics the pre-delegation behaviour for
+/// single-shard tests that don't spin up a coordinator.
+#[must_use]
+pub fn install_replicas_locally(
+ bus: Rc<IggyMessageBus>,
+ on_message: MessageHandler,
+) -> AcceptedReplicaFn {
+ Rc::new(move |stream, peer_id| {
+ installer::install_replica_stream(&bus, peer_id, stream, on_message.clone());
+ })
+}
+
+/// Build an [`AcceptedClientFn`] that mints a local client id (top 16 bits =
+/// `shard_id`, bottom 112 bits = per-call counter) and installs the client
+/// stream directly on the given bus. Tests use this to bypass the shard-0
+/// coordinator while keeping the same install plumbing the production path
+/// exercises.
+#[must_use]
+pub fn install_clients_locally(
+ bus: Rc<IggyMessageBus>,
+ on_request: RequestHandler,
+) -> AcceptedClientFn {
+ let counter: Rc<Cell<u128>> = Rc::new(Cell::new(1));
+ let shard_id = u128::from(bus.shard_id());
+ Rc::new(move |stream| {
+ let seq = counter.get();
+ counter.set(seq.wrapping_add(1));
+ let client_id = (shard_id << 112) | seq;
+ installer::install_client_stream(&bus, client_id, stream, on_request.clone());
+ })
+}
diff --git a/core/message_bus/tests/connection_lost_notify.rs b/core/message_bus/tests/connection_lost_notify.rs
new file mode 100644
index 000000000..8babea3ad
--- /dev/null
+++ b/core/message_bus/tests/connection_lost_notify.rs
@@ -0,0 +1,122 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Both the writer task and the reader task observe an abnormal
+//! connection close. They MUST emit exactly one `connection_lost`
+//! notification per peer disconnect so shard 0 doesn't churn the
+//! mapping broadcast twice in a row.
+
+mod common;
+
+use common::{install_replicas_locally, loopback};
+use message_bus::IggyMessageBus;
+use message_bus::connector::{DEFAULT_RECONNECT_PERIOD, start as start_connector};
+use message_bus::replica_listener::{MessageHandler, bind, run};
+use std::cell::Cell;
+use std::rc::Rc;
+use std::time::Duration;
+
+const CLUSTER: u128 = 0xFEED;
+
+#[compio::test]
+async fn connection_lost_fires_exactly_once_per_peer_disconnect() {
+ // bus0 dials bus1 (directional rule: lower id dials higher id).
+ let bus0 = Rc::new(IggyMessageBus::new(0));
+ let bus1 = Rc::new(IggyMessageBus::new(0));
+
+ let counter: Rc<Cell<u32>> = Rc::new(Cell::new(0));
+ let notified_peer: Rc<Cell<Option<u8>>> = Rc::new(Cell::new(None));
+ {
+ let counter = Rc::clone(&counter);
+ let notified_peer = Rc::clone(¬ified_peer);
+ bus0.set_connection_lost_fn(Rc::new(move |peer_id| {
+ counter.set(counter.get() + 1);
+ notified_peer.set(Some(peer_id));
+ }));
+ }
+
+ let on_message: MessageHandler = Rc::new(|_, _| {});
+
+ // bus1 listens; bus0 dials.
+ let (l1, addr1) = bind(loopback()).await.unwrap();
+ let token_for_l1 = bus1.token();
+ let accept_1 = install_replicas_locally(bus1.clone(), on_message.clone());
+ let l1_handle = compio::runtime::spawn(async move {
+ run(
+ l1,
+ token_for_l1,
+ CLUSTER,
+ 1,
+ 2,
+ accept_1,
+ message_bus::framing::MAX_MESSAGE_SIZE,
+ )
+ .await;
+ });
+ bus1.track_background(l1_handle);
+
+ let dial_0 = install_replicas_locally(bus0.clone(), on_message.clone());
+ start_connector(
+ &bus0,
+ CLUSTER,
+ 0,
+ vec![(1u8, addr1)],
+ dial_0,
+ DEFAULT_RECONNECT_PERIOD,
+ )
+ .await;
+
+ let deadline = std::time::Instant::now() + Duration::from_secs(2);
+ while !bus0.replicas().contains(1) {
+ assert!(
+ std::time::Instant::now() < deadline,
+ "bus0 failed to establish connection to bus1"
+ );
+ compio::time::sleep(Duration::from_millis(10)).await;
+ }
+
+ // Drop bus1 entirely: this closes the listener and every accepted
+ // connection's fd, so bus0's reader sees EOF and bus0's writer sees
+ // a broken pipe on its next flush / close notification.
+ bus1.shutdown(Duration::from_secs(2)).await;
+ drop(bus1);
+
+ // Wait for bus0 to observe the loss and eject the peer.
+ let deadline = std::time::Instant::now() + Duration::from_secs(2);
+ while bus0.replicas().contains(1) {
+ assert!(
+ std::time::Instant::now() < deadline,
+ "bus0 never observed the disconnect"
+ );
+ compio::time::sleep(Duration::from_millis(10)).await;
+ }
+
+ // Small grace window: the second half (reader or writer) may still
+ // be winding down. If the guard works, it will run through the
+ // post-exit block and skip the notify.
+ compio::time::sleep(Duration::from_millis(100)).await;
+
+ assert_eq!(
+ counter.get(),
+ 1,
+ "notify_connection_lost must fire exactly once per peer disconnect, got {}",
+ counter.get(),
+ );
+ assert_eq!(notified_peer.get(), Some(1u8));
+
+ bus0.shutdown(Duration::from_secs(2)).await;
+}
diff --git a/core/message_bus/tests/directional_connection.rs b/core/message_bus/tests/directional_connection.rs
new file mode 100644
index 000000000..ac8a3a7d1
--- /dev/null
+++ b/core/message_bus/tests/directional_connection.rs
@@ -0,0 +1,116 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Both replicas spin up their inbound listener AND their connector at the
+//! same time. The directional rule (lower id dials, higher id accepts)
+//! must yield exactly one TCP connection between them, regardless of
+//! scheduling.
+
+mod common;
+
+use common::{install_replicas_locally, loopback};
+use message_bus::IggyMessageBus;
+use message_bus::connector::{DEFAULT_RECONNECT_PERIOD, start as start_connector};
+use message_bus::replica_listener::{MessageHandler, bind, run};
+use std::rc::Rc;
+use std::time::Duration;
+
+const CLUSTER: u128 = 0xBEEF;
+
+#[compio::test]
+async fn lower_id_dials_higher_id_accepts() {
+ let bus0 = Rc::new(IggyMessageBus::new(0));
+ let bus1 = Rc::new(IggyMessageBus::new(0));
+
+ let on_message: MessageHandler = Rc::new(|_, _| {});
+
+ // Both buses bind their inbound listener.
+ let (l0, addr0) = bind(loopback()).await.unwrap();
+ let (l1, addr1) = bind(loopback()).await.unwrap();
+
+ let token_for_l0 = bus0.token();
+ let accept_0 = install_replicas_locally(bus0.clone(), on_message.clone());
+ let l0_handle = compio::runtime::spawn(async move {
+ run(
+ l0,
+ token_for_l0,
+ CLUSTER,
+ 0,
+ 2,
+ accept_0,
+ message_bus::framing::MAX_MESSAGE_SIZE,
+ )
+ .await;
+ });
+ bus0.track_background(l0_handle);
+
+ let token_for_l1 = bus1.token();
+ let accept_1 = install_replicas_locally(bus1.clone(), on_message.clone());
+ let l1_handle = compio::runtime::spawn(async move {
+ run(
+ l1,
+ token_for_l1,
+ CLUSTER,
+ 1,
+ 2,
+ accept_1,
+ message_bus::framing::MAX_MESSAGE_SIZE,
+ )
+ .await;
+ });
+ bus1.track_background(l1_handle);
+
+ // Both buses start their connector with the full peer list.
+ let peers = vec![(0u8, addr0), (1u8, addr1)];
+
+ let dial_0 = install_replicas_locally(bus0.clone(), on_message.clone());
+ start_connector(
+ &bus0,
+ CLUSTER,
+ 0,
+ peers.clone(),
+ dial_0,
+ DEFAULT_RECONNECT_PERIOD,
+ )
+ .await;
+ let dial_1 = install_replicas_locally(bus1.clone(), on_message.clone());
+ start_connector(&bus1, CLUSTER, 1, peers, dial_1, DEFAULT_RECONNECT_PERIOD).await;
+
+ // Wait for the directional connection to settle.
+ let deadline = std::time::Instant::now() + Duration::from_secs(2);
+ while !(bus0.replicas().contains(1) && bus1.replicas().contains(0)) {
+ assert!(
+ std::time::Instant::now() < deadline,
+ "replicas did not connect: bus0={} bus1={}",
+ bus0.replicas().len(),
+ bus1.replicas().len()
+ );
+ compio::time::sleep(Duration::from_millis(10)).await;
+ }
+
+ // Each side has exactly one peer connection.
+ assert_eq!(bus0.replicas().len(), 1);
+ assert_eq!(bus1.replicas().len(), 1);
+ // bus0 must NOT have dialed itself, and must have peer 1.
+ assert!(bus0.replicas().contains(1));
+ // bus1 must NOT have dialed peer 0 (directional rule); it has the
+ // inbound entry for 0.
+ assert!(bus1.replicas().contains(0));
+
+ bus0.shutdown(Duration::from_secs(2)).await;
+ bus1.shutdown(Duration::from_secs(2)).await;
+}
diff --git a/core/message_bus/tests/duplicate_client_id.rs b/core/message_bus/tests/duplicate_client_id.rs
new file mode 100644
index 000000000..2335e4324
--- /dev/null
+++ b/core/message_bus/tests/duplicate_client_id.rs
@@ -0,0 +1,157 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Shard 0 mints monotonic client ids so the registry should never see a
+//! collision. If one leaks in anyway (bad foreign mint, wrap at 2^112),
+//! the installer must drop the duplicate fd instead of panicking. This
+//! test forces the collision and verifies the first entry survives while
+//! the second is dropped cleanly.
+
+mod common;
+
+use common::{header_only, loopback};
+use compio::net::{TcpListener, TcpStream};
+use iggy_binary_protocol::Command2;
+use message_bus::IggyMessageBus;
+use message_bus::client_listener::RequestHandler;
+use message_bus::framing::write_message;
+use message_bus::installer::install_client_stream;
+use std::cell::Cell;
+use std::rc::Rc;
+use std::time::Duration;
+
+#[allow(clippy::future_not_send)]
+async fn loopback_pair() -> TcpStream {
+ let listener = TcpListener::bind(loopback()).await.expect("bind");
+ let addr = listener.local_addr().expect("local_addr");
+ let accept = compio::runtime::spawn(async move {
+ let (stream, _) = listener.accept().await.expect("accept");
+ stream
+ });
+ let client = TcpStream::connect(addr).await.expect("connect");
+ let _server = accept.await.expect("accept task");
+ client
+}
+
+#[compio::test]
+#[allow(clippy::future_not_send)]
+async fn duplicate_install_drops_second_fd_without_panic() {
+ let bus = Rc::new(IggyMessageBus::new(0));
+ let on_request: RequestHandler = Rc::new(|_, _| {});
+
+ let client_id: u128 = 0x42;
+ let first = loopback_pair().await;
+ install_client_stream(&bus, client_id, first, on_request.clone());
+ assert!(bus.clients().contains(client_id));
+ assert_eq!(bus.clients().len(), 1);
+
+ // Collision: same client_id, second stream. Must not panic, must not
+ // double-register.
+ let second = loopback_pair().await;
+ install_client_stream(&bus, client_id, second, on_request);
+ assert!(bus.clients().contains(client_id));
+ assert_eq!(
+ bus.clients().len(),
+ 1,
+ "duplicate install must not inflate the registry"
+ );
+
+ let outcome = bus.shutdown(Duration::from_secs(2)).await;
+ assert_eq!(
+ outcome.force, 0,
+ "duplicate's writer task must exit cleanly via dropped sender"
+ );
+}
+
+#[allow(clippy::future_not_send)]
+async fn loopback_peer_pair() -> (TcpStream, TcpStream) {
+ let listener = TcpListener::bind(loopback()).await.expect("bind");
+ let addr = listener.local_addr().expect("local_addr");
+ let accept = compio::runtime::spawn(async move {
+ let (stream, _) = listener.accept().await.expect("accept");
+ stream
+ });
+ let client = TcpStream::connect(addr).await.expect("connect");
+ let server = accept.await.expect("accept task");
+ (client, server)
+}
+
+#[compio::test]
+#[allow(clippy::future_not_send)]
+async fn orphan_reader_from_losing_install_does_not_invoke_on_request() {
+ // An orphan reader spawned by the losing side of a duplicate-client-id
+ // install must NOT route inbound requests via `on_request`: the
+ // registry entry that `client_id` keys is the winner's, so anything
+ // forwarded through the winner's entry would land on the wrong
+ // consensus path. The aborted flag set in `install_client_stream` on
+ // `AlreadyRegistered` stops the read loop before it dispatches.
+ let bus = Rc::new(IggyMessageBus::new(0));
+
+ let on_request_count: Rc<Cell<u32>> = Rc::new(Cell::new(0));
+ let on_request: RequestHandler = {
+ let c = Rc::clone(&on_request_count);
+ Rc::new(move |_, _| {
+ c.set(c.get() + 1);
+ })
+ };
+
+ let client_id: u128 = 0xdead_beef;
+
+ // First install wins the slot.
+ let (first_local, _first_peer) = loopback_peer_pair().await;
+ install_client_stream(&bus, client_id, first_local, on_request.clone());
+ assert!(bus.clients().contains(client_id));
+
+ // Second install loses the `AlreadyRegistered` race. Its spawned
+ // reader is holding `second_local`'s read half; we drive that reader
+ // by writing a `Request` frame from `second_peer`.
+ let (second_local, mut second_peer) = loopback_peer_pair().await;
+ install_client_stream(&bus, client_id, second_local, on_request);
+
+ // Push a valid Request frame down the orphan reader's wire. Before
+ // the fix, the orphan would call `on_request(client_id, msg)` and
+ // bump the counter; with the fix the `aborted` guard skips it.
+ let msg = header_only(Command2::Request, 0, 0);
+ write_message(&mut second_peer, msg)
+ .await
+ .expect("write Request");
+ // Give the orphan reader a tick to observe the bytes and hit the
+ // aborted check.
+ compio::time::sleep(Duration::from_millis(50)).await;
+
+ assert_eq!(
+ on_request_count.get(),
+ 0,
+ "aborted orphan reader must not invoke on_request"
+ );
+ assert!(
+ bus.clients().contains(client_id),
+ "winner entry must still be present"
+ );
+
+ // Close the orphan's peer side so its reader observes EOF and runs
+ // its post-loop block. With the fix, it must NOT evict the winner.
+ drop(second_peer);
+ compio::time::sleep(Duration::from_millis(50)).await;
+ assert!(
+ bus.clients().contains(client_id),
+ "aborted orphan reader must not close_peer on the winner"
+ );
+
+ let outcome = bus.shutdown(Duration::from_secs(2)).await;
+ assert_eq!(outcome.force, 0);
+}
diff --git a/core/message_bus/tests/graceful_shutdown.rs b/core/message_bus/tests/graceful_shutdown.rs
new file mode 100644
index 000000000..552fb97b9
--- /dev/null
+++ b/core/message_bus/tests/graceful_shutdown.rs
@@ -0,0 +1,134 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Open several real client connections, call `bus.shutdown(timeout)`, and
+//! verify that all read tasks exit cleanly within the timeout.
+
+mod common;
+
+use common::{install_clients_locally, loopback};
+use compio::net::TcpStream;
+use message_bus::client_listener::{RequestHandler, bind, run};
+use message_bus::{IggyMessageBus, MessageBus, SendError};
+use std::rc::Rc;
+use std::time::Duration;
+
+#[compio::test]
+async fn drains_all_clients_within_timeout() {
+ let bus = Rc::new(IggyMessageBus::new(0));
+ let on_request: RequestHandler = Rc::new(|_, _| {});
+ let (listener, addr) = bind(loopback()).await.unwrap();
+
+ let token = bus.token();
+ let accept_delegate = install_clients_locally(bus.clone(), on_request);
+ let accept = compio::runtime::spawn(async move {
+ run(listener, token, accept_delegate).await;
+ });
+ bus.track_background(accept);
+
+ // Open N clients. The streams are kept alive in `clients` so the
+ // listener observes 5 live connections; we never read from them.
+ #[allow(clippy::collection_is_never_read)]
+ let mut clients: Vec<TcpStream> = Vec::with_capacity(5);
+ for _ in 0..5 {
+ let stream = TcpStream::connect(addr).await.unwrap();
+ clients.push(stream);
+ }
+
+ // Wait for the listener to register all of them.
+ let deadline = std::time::Instant::now() + Duration::from_secs(2);
+ while bus.clients().len() < 5 {
+ assert!(
+ std::time::Instant::now() < deadline,
+ "listener never registered all clients"
+ );
+ compio::time::sleep(Duration::from_millis(5)).await;
+ }
+
+ let outcome = bus.shutdown(Duration::from_secs(2)).await;
+ assert_eq!(outcome.force, 0, "no connection should need force-cancel");
+ // Each connection now has two tasks (writer + reader) so 5 conns -> 10.
+ assert_eq!(outcome.clean, 10);
+ assert!(bus.clients().is_empty());
+ assert!(bus.is_shutting_down());
+
+ // Sends after shutdown must fail with the right variant.
+ let dummy = common::header_only(iggy_binary_protocol::Command2::Reply, 0, 0);
+ let err = bus
+ .send_to_client(1u128 << 112, dummy.into_frozen())
+ .await
+ .unwrap_err();
+ assert!(matches!(err, SendError::BusShuttingDown));
+}
+
+/// Regression: a misbehaving background task that ignores the shutdown
+/// token and outlives the total budget must not starve connection drain.
+/// Connection drain runs BEFORE background drain so writer tasks still
+/// flush cleanly when a rogue background task exists.
+#[compio::test]
+async fn connection_drain_precedes_slow_background() {
+ let bus = Rc::new(IggyMessageBus::new(0));
+ let on_request: RequestHandler = Rc::new(|_, _| {});
+ let (listener, addr) = bind(loopback()).await.unwrap();
+
+ let token = bus.token();
+ let accept_delegate = install_clients_locally(bus.clone(), on_request);
+ let accept = compio::runtime::spawn(async move {
+ run(listener, token, accept_delegate).await;
+ });
+ bus.track_background(accept);
+
+ // Rogue background task: sleeps 3s without observing the shutdown
+ // token. Under the old ordering this would consume the entire
+ // 1s shutdown budget before connection drain starts. Under the new
+ // ordering connections drain first and the rogue task is forced.
+ let rogue = compio::runtime::spawn(async {
+ compio::time::sleep(Duration::from_secs(3)).await;
+ });
+ bus.track_background(rogue);
+
+ #[allow(clippy::collection_is_never_read)]
+ let mut clients: Vec<TcpStream> = Vec::with_capacity(2);
+ for _ in 0..2 {
+ clients.push(TcpStream::connect(addr).await.unwrap());
+ }
+
+ let deadline = std::time::Instant::now() + Duration::from_secs(2);
+ while bus.clients().len() < 2 {
+ assert!(
+ std::time::Instant::now() < deadline,
+ "listener never registered clients"
+ );
+ compio::time::sleep(Duration::from_millis(5)).await;
+ }
+
+ let outcome = bus.shutdown(Duration::from_secs(1)).await;
+ // Clients still drained cleanly: writer + reader per connection.
+ assert_eq!(
+ outcome.force, 0,
+ "connection drain must not be forced by rogue background task"
+ );
+ assert_eq!(outcome.clean, 4, "2 clients * (writer + reader)");
+ // The rogue background task (plus the accept loop, which exits on
+ // the shutdown token) is accounted for in background counters. At
+ // least one background task was forced.
+ assert!(
+ outcome.background_force >= 1,
+ "rogue background task should be force-cancelled, got {outcome:?}"
+ );
+ assert!(bus.clients().is_empty());
+}
diff --git a/core/message_bus/tests/head_of_line.rs b/core/message_bus/tests/head_of_line.rs
new file mode 100644
index 000000000..031c3305c
--- /dev/null
+++ b/core/message_bus/tests/head_of_line.rs
@@ -0,0 +1,151 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Verify that a slow peer cannot stall sends to other peers.
+//!
+//! Setup: a sender bus has two replica connections. Peer A reads normally;
+//! peer B's listener accepts but then drops the read half (so the sender's
+//! kernel TCP send buffer eventually fills, blocking peer B's writer task).
+//! With the per-peer queue model, sends to peer A must remain O(microseconds)
+//! regardless of how blocked peer B is.
+
+mod common;
+
+use common::{header_only, install_replicas_locally, loopback};
+use iggy_binary_protocol::Command2;
+use message_bus::connector::{DEFAULT_RECONNECT_PERIOD, start as start_connector};
+use message_bus::replica_listener::{MessageHandler, bind, run};
+use message_bus::{IggyMessageBus, MessageBus};
+use std::cell::Cell;
+use std::rc::Rc;
+use std::time::{Duration, Instant};
+
+const CLUSTER: u128 = 0xF00D;
+
+#[compio::test]
+async fn slow_peer_does_not_block_other_peers() {
+ // Receiving bus 1 (peer A): drains messages normally.
+ let bus_a = Rc::new(IggyMessageBus::new(0));
+ let received_a = Rc::new(Cell::new(0usize));
+ let received_a_clone = received_a.clone();
+ let on_a: MessageHandler = Rc::new(move |_, _| {
+ received_a_clone.set(received_a_clone.get() + 1);
+ });
+ let (la, addr_a) = bind(loopback()).await.unwrap();
+ let token_a = bus_a.token();
+ let accept_a = install_replicas_locally(bus_a.clone(), on_a.clone());
+ let la_handle = compio::runtime::spawn(async move {
+ run(
+ la,
+ token_a,
+ CLUSTER,
+ 1,
+ 3,
+ accept_a,
+ message_bus::framing::MAX_MESSAGE_SIZE,
+ )
+ .await;
+ });
+ bus_a.track_background(la_handle);
+
+ // Receiving bus 2 (peer B): handler is a no-op, but the listener stays
+ // up. The kernel TCP recv buffer is small, so once we send enough to
+ // peer B without it draining, peer B's writer task on the sender side
+ // will eventually block on writev.
+ let bus_b = Rc::new(IggyMessageBus::new(0));
+ let on_b: MessageHandler = Rc::new(|_, _| {});
+ let (lb, addr_b) = bind(loopback()).await.unwrap();
+ let token_b = bus_b.token();
+ let accept_b = install_replicas_locally(bus_b.clone(), on_b);
+ let lb_handle = compio::runtime::spawn(async move {
+ run(
+ lb,
+ token_b,
+ CLUSTER,
+ 2,
+ 3,
+ accept_b,
+ message_bus::framing::MAX_MESSAGE_SIZE,
+ )
+ .await;
+ });
+ bus_b.track_background(lb_handle);
+
+ // Sender bus 0 dials both A and B.
+ let bus0 = Rc::new(IggyMessageBus::with_capacity(0, 16));
+ let dial_0: MessageHandler = Rc::new(|_, _| {});
+ let dial_delegate = install_replicas_locally(bus0.clone(), dial_0);
+ start_connector(
+ &bus0,
+ CLUSTER,
+ 0,
+ vec![(1, addr_a), (2, addr_b)],
+ dial_delegate,
+ DEFAULT_RECONNECT_PERIOD,
+ )
+ .await;
+
+ let deadline = Instant::now() + Duration::from_secs(2);
+ while !(bus0.replicas().contains(1) && bus0.replicas().contains(2)) {
+ assert!(Instant::now() < deadline, "both replicas must connect");
+ compio::time::sleep(Duration::from_millis(5)).await;
+ }
+
+ // Time a single send_to_replica to the FAST peer A. With per-peer
+ // queues this must complete in microseconds even if peer B is busy.
+ let send_a_start = Instant::now();
+ bus0.send_to_replica(1, header_only(Command2::Prepare, 0, 0).into_frozen())
+ .await
+ .expect("send to A");
+ let send_a_elapsed = send_a_start.elapsed();
+ assert!(
+ send_a_elapsed < Duration::from_millis(50),
+ "send_to_replica(A) took {send_a_elapsed:?}; head-of-line blocking?"
+ );
+
+ // Even after issuing many sends to B (which will eventually pile up in
+ // the kernel buffer and the per-peer queue), sends to A must stay fast.
+ for _ in 0..50 {
+ let _ = bus0
+ .send_to_replica(2, header_only(Command2::Prepare, 0, 0).into_frozen())
+ .await;
+ }
+ let send_a2_start = Instant::now();
+ bus0.send_to_replica(1, header_only(Command2::Prepare, 0, 0).into_frozen())
+ .await
+ .expect("send to A after B saturation");
+ let send_a2_elapsed = send_a2_start.elapsed();
+ assert!(
+ send_a2_elapsed < Duration::from_millis(50),
+ "second send_to_replica(A) took {send_a2_elapsed:?}; HOL after B saturation"
+ );
+
+ // Wait for peer A to receive at least our two messages.
+ let deadline = Instant::now() + Duration::from_secs(2);
+ while received_a.get() < 2 {
+ assert!(
+ Instant::now() < deadline,
+ "peer A received only {}/2",
+ received_a.get()
+ );
+ compio::time::sleep(Duration::from_millis(5)).await;
+ }
+
+ bus0.shutdown(Duration::from_secs(2)).await;
+ bus_a.shutdown(Duration::from_secs(2)).await;
+ bus_b.shutdown(Duration::from_secs(2)).await;
+}
diff --git a/core/message_bus/tests/reconnect.rs b/core/message_bus/tests/reconnect.rs
new file mode 100644
index 000000000..0eef98e94
--- /dev/null
+++ b/core/message_bus/tests/reconnect.rs
@@ -0,0 +1,80 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! The dialer's connector keeps retrying after an initial connect failure.
+//! Once the listener comes up on the same port, the connector picks it up
+//! within one reconnect period.
+
+mod common;
+
+use common::{install_replicas_locally, loopback};
+use message_bus::IggyMessageBus;
+use message_bus::connector::start as start_connector;
+use message_bus::replica_listener::{MessageHandler, bind, run};
+use std::rc::Rc;
+use std::time::Duration;
+
+const CLUSTER: u128 = 0xABCD;
+
+#[compio::test]
+async fn periodic_retry_picks_up_late_listener() {
+ // Reserve a free port by binding then dropping.
+ let (probe, addr) = bind(loopback()).await.unwrap();
+ drop(probe);
+
+ // Bus 0 starts its connector targeting that addr - nothing is listening
+ // yet, so the first attempt fails. The reconnect period is short so the
+ // test does not have to wait long.
+ let bus0 = Rc::new(IggyMessageBus::new(0));
+ let on_message: MessageHandler = Rc::new(|_, _| {});
+ let period = Duration::from_millis(100);
+ let dial_delegate = install_replicas_locally(bus0.clone(), on_message.clone());
+ start_connector(&bus0, CLUSTER, 0, vec![(1, addr)], dial_delegate, period).await;
+ assert!(!bus0.replicas().contains(1), "first connect should fail");
+
+ // Bring bus 1 online on the same address.
+ let bus1 = Rc::new(IggyMessageBus::new(0));
+ let (l1, _) = bind(addr).await.unwrap();
+ let token_for_l1 = bus1.token();
+ let accept_delegate = install_replicas_locally(bus1.clone(), on_message.clone());
+ let l1_handle = compio::runtime::spawn(async move {
+ run(
+ l1,
+ token_for_l1,
+ CLUSTER,
+ 1,
+ 2,
+ accept_delegate,
+ message_bus::framing::MAX_MESSAGE_SIZE,
+ )
+ .await;
+ });
+ bus1.track_background(l1_handle);
+
+ // Within one or two reconnect periods, bus 0 should connect.
+ let deadline = std::time::Instant::now() + Duration::from_secs(2);
+ while !bus0.replicas().contains(1) {
+ assert!(
+ std::time::Instant::now() < deadline,
+ "reconnect periodic did not establish the connection"
+ );
+ compio::time::sleep(Duration::from_millis(20)).await;
+ }
+
+ bus0.shutdown(Duration::from_secs(2)).await;
+ bus1.shutdown(Duration::from_secs(2)).await;
+}
diff --git a/core/message_bus/tests/reconnect_skip_connected.rs b/core/message_bus/tests/reconnect_skip_connected.rs
new file mode 100644
index 000000000..5d9a77b3b
--- /dev/null
+++ b/core/message_bus/tests/reconnect_skip_connected.rs
@@ -0,0 +1,110 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Once a peer is registered on the cluster the periodic reconnect
+//! sweep must NOT dial it again. A redial would produce a second TCP
+//! connection that either races the live one or flaps the mapping
+//! across shards.
+
+mod common;
+
+use common::{install_replicas_locally, loopback};
+use message_bus::IggyMessageBus;
+use message_bus::connector::start as start_connector;
+use message_bus::replica_listener::{MessageHandler, bind, run};
+use std::rc::Rc;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU32, Ordering};
+use std::time::Duration;
+
+const CLUSTER: u128 = 0xC2C2;
+
+#[compio::test]
+async fn periodic_reconnect_skips_already_connected_peer() {
+ let accept_count = Arc::new(AtomicU32::new(0));
+ let bus1 = Rc::new(IggyMessageBus::new(0));
+ let bus0 = Rc::new(IggyMessageBus::new(0));
+
+ // Listener on bus1: count accepted TCP connections via a wrapper
+ // around the install delegate.
+ let (l1, addr1) = bind(loopback()).await.unwrap();
+ let on_message: MessageHandler = Rc::new(|_, _| {});
+ let accept_inner = install_replicas_locally(bus1.clone(), on_message.clone());
+ let accept_counter = Arc::clone(&accept_count);
+ let accept_delegate: message_bus::AcceptedReplicaFn = Rc::new(move |stream, peer_id| {
+ accept_counter.fetch_add(1, Ordering::SeqCst);
+ accept_inner(stream, peer_id);
+ });
+ let token_for_l1 = bus1.token();
+ let l1_handle = compio::runtime::spawn(async move {
+ run(
+ l1,
+ token_for_l1,
+ CLUSTER,
+ 1,
+ 2,
+ accept_delegate,
+ message_bus::framing::MAX_MESSAGE_SIZE,
+ )
+ .await;
+ });
+ bus1.track_background(l1_handle);
+
+ // bus0 connector: very short reconnect period so the sweep fires
+ // multiple times within the test window. If the skip guard is
+ // missing, bus1's listener will see N extra accepts.
+ let period = Duration::from_millis(50);
+ let dial_delegate = install_replicas_locally(bus0.clone(), on_message.clone());
+ start_connector(&bus0, CLUSTER, 0, vec![(1u8, addr1)], dial_delegate, period).await;
+
+ // Wait for the initial connection to settle on both sides.
+ let deadline = std::time::Instant::now() + Duration::from_secs(2);
+ loop {
+ if bus0.replicas().contains(1) && accept_count.load(Ordering::SeqCst) >= 1 {
+ break;
+ }
+ assert!(
+ std::time::Instant::now() < deadline,
+ "initial connect failed: bus0.contains(1)={}, accepts={}",
+ bus0.replicas().contains(1),
+ accept_count.load(Ordering::SeqCst),
+ );
+ compio::time::sleep(Duration::from_millis(10)).await;
+ }
+ assert_eq!(
+ accept_count.load(Ordering::SeqCst),
+ 1,
+ "initial connect should produce exactly one accept"
+ );
+
+ // Let the periodic sweep fire at least ~5 times.
+ compio::time::sleep(period * 6).await;
+
+ // The connection is still live AND no extra accepts happened.
+ assert!(
+ bus0.replicas().contains(1),
+ "peer should still be connected"
+ );
+ assert_eq!(
+ accept_count.load(Ordering::SeqCst),
+ 1,
+ "periodic sweep must not redial a live peer"
+ );
+
+ bus0.shutdown(Duration::from_secs(2)).await;
+ bus1.shutdown(Duration::from_secs(2)).await;
+}
diff --git a/core/message_bus/tests/replica_roundtrip.rs b/core/message_bus/tests/replica_roundtrip.rs
new file mode 100644
index 000000000..e557b501a
--- /dev/null
+++ b/core/message_bus/tests/replica_roundtrip.rs
@@ -0,0 +1,127 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Two real `IggyMessageBus` instances act as replica 0 and replica 1.
+//! Replica 0 dials replica 1; the Ping handshake completes; replica 0 sends
+//! a Prepare; replica 1's `on_message` callback receives it; replica 1
+//! responds with `PrepareOk` over its own bus; replica 0 receives it.
+
+mod common;
+
+use async_channel::Receiver;
+use common::{header_only, install_replicas_locally, loopback};
+use iggy_binary_protocol::Command2;
+use message_bus::connector::{DEFAULT_RECONNECT_PERIOD, start as start_connector};
+use message_bus::replica_listener::{MessageHandler, bind, run};
+use message_bus::{IggyMessageBus, MessageBus};
+use std::rc::Rc;
+use std::time::Duration;
+
+const CLUSTER: u128 = 0xCAFE;
+
+#[compio::test]
+async fn two_replicas_exchange_prepare_and_ack() {
+ // Bring up replica 1's inbound listener first so replica 0's connect succeeds.
+ let bus1 = Rc::new(IggyMessageBus::new(0));
+ let (tx1, rx1) = async_channel::bounded::<u8>(8);
+ let on_message1: MessageHandler = Rc::new(move |peer, msg| {
+ tx1.try_send(msg.header().command as u8).ok();
+ let _ = peer;
+ });
+ let (l1, addr1) = bind(loopback()).await.expect("bind r1");
+
+ let token_for_l1 = bus1.token();
+ let accept_delegate_1 = install_replicas_locally(bus1.clone(), on_message1.clone());
+ let l1_handle = compio::runtime::spawn(async move {
+ run(
+ l1,
+ token_for_l1,
+ CLUSTER,
+ 1,
+ 2,
+ accept_delegate_1,
+ message_bus::framing::MAX_MESSAGE_SIZE,
+ )
+ .await;
+ });
+ bus1.track_background(l1_handle);
+
+ // Replica 0: outbound connector dials replica 1.
+ let bus0 = Rc::new(IggyMessageBus::new(0));
+ let (tx0, rx0) = async_channel::bounded::<u8>(8);
+ let on_message0: MessageHandler = Rc::new(move |peer, msg| {
+ tx0.try_send(msg.header().command as u8).ok();
+ let _ = peer;
+ });
+ let dial_delegate_0 = install_replicas_locally(bus0.clone(), on_message0);
+
+ start_connector(
+ &bus0,
+ CLUSTER,
+ 0,
+ vec![(1, addr1)],
+ dial_delegate_0,
+ DEFAULT_RECONNECT_PERIOD,
+ )
+ .await;
+
+ // Wait until replica 0 has the connection registered.
+ wait_until(|| bus0.replicas().contains(1), Duration::from_secs(2)).await;
+
+ // Replica 0 sends a Prepare to replica 1.
+ let prepare = header_only(Command2::Prepare, CLUSTER, 0);
+ bus0.send_to_replica(1, prepare.into_frozen())
+ .await
+ .expect("send prepare");
+
+ // Replica 1 should observe the Prepare via its on_message callback.
+ let cmd = expect_recv(&rx1, Duration::from_secs(2)).await;
+ assert_eq!(cmd, Command2::Prepare as u8);
+
+ // Replica 1 responds with PrepareOk.
+ let ack = header_only(Command2::PrepareOk, CLUSTER, 1);
+ bus1.send_to_replica(0, ack.into_frozen())
+ .await
+ .expect("send ack");
+
+ // Replica 0 should observe the PrepareOk on its outbound read loop.
+ let cmd = expect_recv(&rx0, Duration::from_secs(2)).await;
+ assert_eq!(cmd, Command2::PrepareOk as u8);
+
+ bus0.shutdown(Duration::from_secs(2)).await;
+ bus1.shutdown(Duration::from_secs(2)).await;
+}
+
+#[allow(clippy::future_not_send)]
+async fn wait_until<F: Fn() -> bool>(cond: F, timeout: Duration) {
+ let deadline = std::time::Instant::now() + timeout;
+ while std::time::Instant::now() < deadline {
+ if cond() {
+ return;
+ }
+ compio::time::sleep(Duration::from_millis(10)).await;
+ }
+ panic!("wait_until: condition not met within {timeout:?}");
+}
+
+#[allow(clippy::future_not_send)]
+async fn expect_recv(rx: &Receiver<u8>, timeout: Duration) -> u8 {
+ compio::time::timeout(timeout, rx.recv())
+ .await
+ .expect("expect_recv: timeout")
+ .expect("expect_recv: channel closed")
+}
diff --git a/core/message_bus/tests/shard_zero_gating.rs b/core/message_bus/tests/shard_zero_gating.rs
new file mode 100644
index 000000000..a65fd84d7
--- /dev/null
+++ b/core/message_bus/tests/shard_zero_gating.rs
@@ -0,0 +1,152 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! `replica_io::start_on_shard_zero` only wires TCP I/O when the bus runs
+//! on shard 0. Non-zero shards keep a working bus instance but never bind
+//! a listener, never dial peers, and never register background tasks.
+
+mod common;
+
+use common::{install_clients_locally, install_replicas_locally, loopback};
+use message_bus::IggyMessageBus;
+use message_bus::client_listener::RequestHandler;
+use message_bus::connector::DEFAULT_RECONNECT_PERIOD;
+use message_bus::replica_io::start_on_shard_zero;
+use message_bus::replica_listener::{MessageHandler, bind, run};
+use std::net::SocketAddr;
+use std::rc::Rc;
+use std::time::Duration;
+
+const CLUSTER: u128 = 0xF00D;
+
+#[compio::test]
+async fn shard_zero_binds_listener_and_starts_connector() {
+ // Peer (replica 1) stands up its own bus + inbound listener so the
+ // directional dial from replica 0 succeeds.
+ let peer_bus = Rc::new(IggyMessageBus::new(0));
+ let peer_handler: MessageHandler = Rc::new(|_, _| {});
+ let (peer_listener, peer_addr) = bind(loopback()).await.expect("bind peer listener");
+ let peer_token = peer_bus.token();
+ let peer_accept = install_replicas_locally(peer_bus.clone(), peer_handler.clone());
+ let peer_listen_handle = compio::runtime::spawn(async move {
+ run(
+ peer_listener,
+ peer_token,
+ CLUSTER,
+ 1,
+ 2,
+ peer_accept,
+ message_bus::framing::MAX_MESSAGE_SIZE,
+ )
+ .await;
+ });
+ peer_bus.track_background(peer_listen_handle);
+
+ // Shard 0 bus under test. For this gating test we wire the
+ // install-locally delegates so shard 0 registers inbound / outbound
+ // peers on its own registry (mimicking a single-shard deployment).
+ let bus_zero = Rc::new(IggyMessageBus::new(0));
+ let on_message: MessageHandler = Rc::new(|_, _| {});
+ let on_request: RequestHandler = Rc::new(|_, _| {});
+ let accepted_replica = install_replicas_locally(bus_zero.clone(), on_message.clone());
+ let accepted_client = install_clients_locally(bus_zero.clone(), on_request);
+
+ let bound = start_on_shard_zero(
+ &bus_zero,
+ loopback(),
+ loopback(),
+ CLUSTER,
+ 0,
+ 2,
+ vec![(1u8, peer_addr)],
+ accepted_replica,
+ accepted_client,
+ DEFAULT_RECONNECT_PERIOD,
+ )
+ .await
+ .expect("start_on_shard_zero must succeed on shard 0");
+
+ let bound = bound.expect("shard 0 must return bound listeners");
+ let replica_addr: SocketAddr = bound.replica;
+ let client_addr: SocketAddr = bound.client;
+ assert_ne!(replica_addr.port(), 0, "replica port must be assigned");
+ assert_ne!(client_addr.port(), 0, "client port must be assigned");
+ assert_ne!(
+ replica_addr.port(),
+ client_addr.port(),
+ "replica and client must bind distinct ports"
+ );
+
+ // Wait for the outbound dial to replica 1 to land.
+ let deadline = std::time::Instant::now() + Duration::from_secs(2);
+ while !bus_zero.replicas().contains(1) {
+ assert!(
+ std::time::Instant::now() < deadline,
+ "shard 0 did not register the outbound connection to peer 1"
+ );
+ compio::time::sleep(Duration::from_millis(10)).await;
+ }
+ assert_eq!(bus_zero.replicas().len(), 1);
+
+ bus_zero.shutdown(Duration::from_secs(2)).await;
+ peer_bus.shutdown(Duration::from_secs(2)).await;
+}
+
+#[compio::test]
+async fn non_zero_shard_skips_io() {
+ let bus_one = Rc::new(IggyMessageBus::new(1));
+ let on_message: MessageHandler = Rc::new(|_, _| {});
+ let on_request: RequestHandler = Rc::new(|_, _| {});
+ let accepted_replica = install_replicas_locally(bus_one.clone(), on_message);
+ let accepted_client = install_clients_locally(bus_one.clone(), on_request);
+
+ let dead_peer: SocketAddr = "127.0.0.1:1".parse().unwrap();
+
+ let outcome = start_on_shard_zero(
+ &bus_one,
+ loopback(),
+ loopback(),
+ CLUSTER,
+ 1,
+ 2,
+ vec![(0u8, dead_peer)],
+ accepted_replica,
+ accepted_client,
+ DEFAULT_RECONNECT_PERIOD,
+ )
+ .await
+ .expect("start_on_shard_zero must succeed on non-zero shard (no-op)");
+
+ assert!(outcome.is_none(), "non-zero shard must not bind a listener");
+ assert_eq!(bus_one.replicas().len(), 0);
+ assert_eq!(bus_one.clients().len(), 0);
+
+ let drained = bus_one.shutdown(Duration::from_millis(100)).await;
+ assert_eq!(drained.clean, 0, "no peer entries should have been drained");
+ assert_eq!(
+ drained.force, 0,
+ "no peer entries should have been force-cancelled"
+ );
+ assert_eq!(
+ drained.background_clean, 0,
+ "helper must not register any background tasks on non-zero shards"
+ );
+ assert_eq!(
+ drained.background_force, 0,
+ "helper must not register any background tasks on non-zero shards"
+ );
+}
diff --git a/core/message_bus/tests/tcp_client_roundtrip.rs b/core/message_bus/tests/tcp_client_roundtrip.rs
new file mode 100644
index 000000000..76002e287
--- /dev/null
+++ b/core/message_bus/tests/tcp_client_roundtrip.rs
@@ -0,0 +1,110 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! End-to-end: a real TCP client connects to the consensus client listener,
+//! sends a Request, the handler echoes a Reply back via `bus.send_to_client`,
+//! the client reads the Reply.
+
+mod common;
+
+use common::{header_only, install_clients_locally, loopback};
+use compio::net::TcpStream;
+use iggy_binary_protocol::Command2;
+use message_bus::client_listener::{RequestHandler, bind, run};
+use message_bus::framing;
+use message_bus::{IggyMessageBus, MessageBus};
+use std::rc::Rc;
+use std::time::Duration;
+
+#[compio::test]
+async fn request_reply_round_trip() {
+ let bus = Rc::new(IggyMessageBus::new(7));
+
+ // Handler echoes a Reply back via send_to_client.
+ let bus_for_handler = bus.clone();
+ let on_request: RequestHandler = Rc::new(move |client_id, msg| {
+ assert_eq!(msg.header().command, Command2::Request);
+ let bus = bus_for_handler.clone();
+ compio::runtime::spawn(async move {
+ let reply = header_only(Command2::Reply, 42, 0);
+ bus.send_to_client(client_id, reply.into_frozen())
+ .await
+ .expect("send_to_client should succeed");
+ })
+ .detach();
+ });
+
+ let (listener, addr) = bind(loopback()).await.expect("bind");
+ let token = bus.token();
+ let accept_delegate = install_clients_locally(bus.clone(), on_request);
+ let accept_handle = compio::runtime::spawn(async move {
+ run(listener, token, accept_delegate).await;
+ });
+ bus.track_background(accept_handle);
+
+ // Dial as a raw TCP client.
+ let mut client = TcpStream::connect(addr).await.expect("connect");
+
+ let request = header_only(Command2::Request, 42, 0);
+ framing::write_message(&mut client, request)
+ .await
+ .expect("client write");
+
+ let reply = framing::read_message(&mut client, framing::MAX_MESSAGE_SIZE)
+ .await
+ .expect("client read");
+ assert_eq!(reply.header().command, Command2::Reply);
+ assert_eq!(reply.header().cluster, 42);
+
+ let outcome = bus.shutdown(Duration::from_secs(2)).await;
+ assert_eq!(
+ outcome.force, 0,
+ "graceful shutdown should not force-cancel"
+ );
+}
+
+#[compio::test]
+async fn unexpected_command_is_ignored() {
+ let bus = Rc::new(IggyMessageBus::new(0));
+
+ // Handler should never be called because only Request is accepted.
+ let (tx, rx) = async_channel::bounded::<()>(1);
+ let on_request: RequestHandler = Rc::new(move |_, _| {
+ let _ = tx.try_send(());
+ });
+
+ let (listener, addr) = bind(loopback()).await.unwrap();
+ let token = bus.token();
+ let accept_delegate = install_clients_locally(bus.clone(), on_request);
+ let accept_handle = compio::runtime::spawn(async move {
+ run(listener, token, accept_delegate).await;
+ });
+ bus.track_background(accept_handle);
+
+ let mut client = TcpStream::connect(addr).await.unwrap();
+ let bogus = header_only(Command2::Ping, 0, 0);
+ framing::write_message(&mut client, bogus).await.unwrap();
+
+ // Give the read loop a chance to observe + ignore it.
+ compio::time::sleep(Duration::from_millis(50)).await;
+ assert!(
+ rx.try_recv().is_err(),
+ "handler must not be called for non-Request commands"
+ );
+
+ bus.shutdown(Duration::from_secs(1)).await;
+}
diff --git a/core/message_bus/tests/vectored_batch.rs b/core/message_bus/tests/vectored_batch.rs
new file mode 100644
index 000000000..c8fc10c6b
--- /dev/null
+++ b/core/message_bus/tests/vectored_batch.rs
@@ -0,0 +1,107 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Verify that the writer task batches pipelined sends into a single
+//! syscall and that all messages arrive in order at the receiver.
+
+mod common;
+
+use common::{header_only, install_replicas_locally, loopback};
+use iggy_binary_protocol::Command2;
+use message_bus::connector::{DEFAULT_RECONNECT_PERIOD, start as start_connector};
+use message_bus::replica_listener::{MessageHandler, bind, run};
+use message_bus::{IggyMessageBus, MessageBus};
+use std::cell::Cell;
+use std::rc::Rc;
+use std::time::Duration;
+
+const CLUSTER: u128 = 0xBA7C;
+const N: usize = 200;
+
+#[compio::test]
+async fn writer_batches_pipelined_sends_in_order() {
+ let bus1 = Rc::new(IggyMessageBus::new(0));
+ let received = Rc::new(Cell::new(0usize));
+ let received_for_handler = received.clone();
+ let on_message: MessageHandler = Rc::new(move |_peer, msg| {
+ // The cluster field carries the message sequence number so we can
+ // assert ordering as well as count.
+ let seq = msg.header().cluster as usize;
+ let prev = received_for_handler.get();
+ assert_eq!(seq, prev, "messages must arrive in order");
+ received_for_handler.set(prev + 1);
+ });
+ let (l1, addr1) = bind(loopback()).await.unwrap();
+ let token_for_l1 = bus1.token();
+ let accept_1 = install_replicas_locally(bus1.clone(), on_message.clone());
+ let l1_handle = compio::runtime::spawn(async move {
+ run(
+ l1,
+ token_for_l1,
+ CLUSTER,
+ 1,
+ 2,
+ accept_1,
+ message_bus::framing::MAX_MESSAGE_SIZE,
+ )
+ .await;
+ });
+ bus1.track_background(l1_handle);
+
+ let bus0 = Rc::new(IggyMessageBus::new(0));
+ let on_reply: MessageHandler = Rc::new(|_, _| {});
+ let dial_0 = install_replicas_locally(bus0.clone(), on_reply);
+ start_connector(
+ &bus0,
+ CLUSTER,
+ 0,
+ vec![(1, addr1)],
+ dial_0,
+ DEFAULT_RECONNECT_PERIOD,
+ )
+ .await;
+
+ // Wait for connect.
+ let deadline = std::time::Instant::now() + Duration::from_secs(2);
+ while !bus0.replicas().contains(1) {
+ assert!(std::time::Instant::now() < deadline, "connect timeout");
+ compio::time::sleep(Duration::from_millis(5)).await;
+ }
+
+ // Pipeline N sends back-to-back without yielding. The writer task
+ // should drain many of them in single writev calls.
+ for i in 0..N {
+ let msg = header_only(Command2::Prepare, i as u128, 0);
+ bus0.send_to_replica(1, msg.into_frozen())
+ .await
+ .expect("send should succeed - queue capacity is 256");
+ }
+
+ // Wait for delivery.
+ let deadline = std::time::Instant::now() + Duration::from_secs(5);
+ while received.get() < N {
+ assert!(
+ std::time::Instant::now() < deadline,
+ "received only {}/{N}",
+ received.get()
+ );
+ compio::time::sleep(Duration::from_millis(5)).await;
+ }
+
+ bus0.shutdown(Duration::from_secs(2)).await;
+ bus1.shutdown(Duration::from_secs(2)).await;
+}
diff --git a/core/metadata/src/impls/metadata.rs b/core/metadata/src/impls/metadata.rs
index c362821ec..a5d4f9b7e 100644
--- a/core/metadata/src/impls/metadata.rs
+++ b/core/metadata/src/impls/metadata.rs
@@ -33,8 +33,12 @@ use message_bus::MessageBus;
use std::path::Path;
use tracing::{debug, error, warn};
-const fn freeze_client_reply(message: Message<GenericHeader>) -> Message<GenericHeader> {
- message
+fn freeze_client_reply(
+ message: Message<GenericHeader>,
+) -> iggy_binary_protocol::consensus::iobuf::Frozen<
+ { iggy_binary_protocol::consensus::MESSAGE_ALIGN },
+> {
+ message.into_frozen()
}
#[derive(Debug, Clone)]
@@ -280,7 +284,7 @@ where
#[allow(clippy::future_not_send)]
impl<B, J, S, M> Plane<VsrConsensus<B>> for IggyMetadata<VsrConsensus<B>, J, S, M>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
J: JournalHandle,
J::Target: Journal<J::Storage, Entry = Message<PrepareHeader>, Header = PrepareHeader>,
M: StateMachine<
@@ -369,7 +373,7 @@ where
#[allow(clippy::cast_possible_truncation)]
let is_old_prepare = fence_old_prepare_by_commit(consensus, &header)
|| journal.handle().header(header.op as usize).is_some();
- let message = if is_old_prepare {
+ if is_old_prepare {
warn!(
target: "iggy.metadata.diag",
plane = "metadata",
@@ -380,10 +384,15 @@ where
operation = ?header.operation,
"received old prepare, skipping replication"
);
- message
} else {
- self.replicate(message).await
- };
+ // TODO(hubcio): is this correct? should we replicate and append,
+ // or append and replicate? Chain-replicating before the local
+ // append means a backup that crashes here forwarded a prepare
+ // it never persisted, violating VSR's tail-ahead-of-head
+ // invariant. The hash-chain fence + view change recovers but
+ // burns a view. Worth an ordering decision issue.
+ self.replicate(&message).await;
+ }
// TODO add assertions for valid state here.
@@ -614,7 +623,7 @@ where
impl<B, P, J, S, M> PlaneIdentity<VsrConsensus<B, P>> for IggyMetadata<VsrConsensus<B, P>, J, S, M>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = PipelineEntry>,
J: JournalHandle,
J::Target: Journal<J::Storage, Entry = Message<PrepareHeader>, Header = PrepareHeader>,
@@ -635,7 +644,7 @@ where
impl<B, P, J, S, M> IggyMetadata<VsrConsensus<B, P>, J, S, M>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = PipelineEntry>,
J: JournalHandle,
J::Target: Journal<J::Storage, Entry = Message<PrepareHeader>, Header = PrepareHeader>,
@@ -651,8 +660,10 @@ where
/// - Primary sends to first backup
/// - Each backup forwards to the next
/// - Stops when we would forward back to primary
+ ///
+ /// Borrows the message so the caller retains ownership for journal append.
#[allow(clippy::future_not_send)]
- async fn replicate(&self, message: Message<PrepareHeader>) -> Message<PrepareHeader> {
+ async fn replicate(&self, message: &Message<PrepareHeader>) {
let consensus = self.consensus.as_ref().unwrap();
let journal = self.journal.as_ref().unwrap();
@@ -666,7 +677,9 @@ where
journal.handle().header(idx).is_none(),
"replicate: must not already have prepare"
);
- replicate_to_next_in_chain(consensus, message).await
+ if let Err(e) = replicate_to_next_in_chain(consensus, message).await {
+ tracing::warn!(op = header.op, error = ?e, "chain replication failed");
+ }
}
// TODO: Implement jump_to_newer_op
diff --git a/core/partitions/src/iggy_partition.rs b/core/partitions/src/iggy_partition.rs
index 729cc2132..14504e2b9 100644
--- a/core/partitions/src/iggy_partition.rs
+++ b/core/partitions/src/iggy_partition.rs
@@ -37,8 +37,8 @@ use consensus::{
request_preflight, send_prepare_ok as send_prepare_ok_common,
};
use iggy_binary_protocol::consensus::iobuf::Frozen;
-use iggy_binary_protocol::{GenericHeader, PrepareOkHeader, RequestHeader};
use iggy_binary_protocol::{Message, Operation, PrepareHeader};
+use iggy_binary_protocol::{PrepareOkHeader, RequestHeader};
use iggy_common::{
ConsumerGroupId, ConsumerGroupOffsets, ConsumerKind, ConsumerOffset, ConsumerOffsets,
IggyByteSize, IggyError, IggyTimestamp, PartitionStats, PollingKind, SegmentStorage,
@@ -571,7 +571,7 @@ where
impl<B> IggyPartition<B>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
{
#[must_use]
fn namespace(&self) -> IggyNamespace {
@@ -683,7 +683,7 @@ where
self.on_replicate(prepare).await;
}
- #[allow(clippy::future_not_send)]
+ #[allow(clippy::future_not_send, clippy::too_many_lines)]
pub async fn on_replicate(&mut self, message: Message<PrepareHeader>) {
self.clear_pending_consumer_offset_commits_if_view_changed();
let header = *message.header();
@@ -713,7 +713,7 @@ where
fence_old_prepare_by_commit(consensus, &header)
|| self.log.journal().inner.header_by_op(header.op).is_some()
};
- let message = if is_old_prepare {
+ if is_old_prepare {
emit_partition_diag(
tracing::Level::WARN,
&PartitionDiagEvent::new(
@@ -723,11 +723,21 @@ where
.with_operation(header.operation)
.with_op(header.op),
);
- message
} else {
let consensus = self.consensus();
- replicate_to_next_in_chain(consensus, message).await
- };
+ if let Err(error) = replicate_to_next_in_chain(consensus, &message).await {
+ emit_partition_diag(
+ tracing::Level::WARN,
+ &PartitionDiagEvent::new(
+ self.diag_ctx(),
+ "failed to replicate prepare to next in chain",
+ )
+ .with_operation(header.operation)
+ .with_op(header.op)
+ .with_error(error.to_string()),
+ );
+ }
+ }
if header.op != current_op + 1 {
emit_partition_diag(
@@ -1133,7 +1143,7 @@ where
);
if send_client_replies {
- let reply_buffers = reply.into_generic();
+ let reply_buffers = reply.into_generic().into_frozen();
emit_sim_event(SimEventKind::ClientReplyEmitted, &event);
if let Err(error) = self
diff --git a/core/partitions/src/iggy_partitions.rs b/core/partitions/src/iggy_partitions.rs
index e8f1582dc..e61e9f72a 100644
--- a/core/partitions/src/iggy_partitions.rs
+++ b/core/partitions/src/iggy_partitions.rs
@@ -21,8 +21,7 @@ use crate::IggyPartition;
use crate::types::PartitionsConfig;
use consensus::{Consensus, Plane, PlaneIdentity, VsrConsensus};
use iggy_binary_protocol::{
- Command2, ConsensusHeader, GenericHeader, Message, PrepareHeader, PrepareOkHeader,
- RequestHeader,
+ Command2, ConsensusHeader, PrepareHeader, PrepareOkHeader, RequestHeader,
};
use iggy_common::sharding::{IggyNamespace, LocalIdx, ShardId};
use message_bus::MessageBus;
@@ -193,7 +192,7 @@ where
impl<B> Plane<VsrConsensus<B>> for IggyPartitions<B>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
{
async fn on_request(&self, message: <VsrConsensus<B> as Consensus>::Message<RequestHeader>) {
let namespace = IggyNamespace::from_raw(message.header().namespace);
@@ -246,7 +245,7 @@ where
impl<B> PlaneIdentity<VsrConsensus<B>> for IggyPartitions<B>
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
{
fn is_applicable<H>(&self, message: &<VsrConsensus<B> as Consensus>::Message<H>) -> bool
where
diff --git a/core/partitions/src/journal.rs b/core/partitions/src/journal.rs
index bfa713d34..891ef0c63 100644
--- a/core/partitions/src/journal.rs
+++ b/core/partitions/src/journal.rs
@@ -542,3 +542,50 @@ fn push_selected_batch_fragments(
*last_matching_offset = Some(selection.last_matching_offset);
*matched_messages += selection.matched_messages;
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use iggy_binary_protocol::{Command2, HEADER_SIZE, Message};
+ use journal::Journal;
+
+ fn build_prepare(op: u64, size: usize) -> Message<PrepareHeader> {
+ Message::<PrepareHeader>::new(size).transmute_header(|_, h: &mut PrepareHeader| {
+ h.command = Command2::Prepare;
+ h.op = op;
+ h.size = u32::try_from(size).expect("size fits in u32");
+ })
+ }
+
+ #[compio::test]
+ async fn entry_round_trips_bytes_for_retransmit() {
+ let journal = PartitionJournal::<PartitionJournalMemStorage>::default();
+
+ let payload_size = HEADER_SIZE + 64;
+ let prepare = build_prepare(3, payload_size);
+ let expected_bytes = prepare.as_slice().to_vec();
+ let frozen = prepare.into_frozen();
+
+ journal.append(frozen).await.expect("append");
+
+ let header = journal.header_by_op(3).expect("header for op 3");
+ let entry = journal
+ .entry(&header)
+ .await
+ .expect("entry for op 3 must exist");
+
+ assert_eq!(
+ entry.as_slice(),
+ expected_bytes.as_slice(),
+ "retransmit path must read back the exact bytes that were appended; \
+ cloning the returned Frozen is the sole payload copy"
+ );
+
+ let cloned = entry.clone();
+ assert_eq!(
+ cloned.as_slice(),
+ entry.as_slice(),
+ "cloning a journal entry must yield identical bytes (refcount bump, not deep copy)"
+ );
+ }
+}
diff --git a/core/server/config.toml b/core/server/config.toml
index 367d5f995..42ed4f2f7 100644
--- a/core/server/config.toml
+++ b/core/server/config.toml
@@ -552,6 +552,8 @@ name = "iggy-cluster"
name = "iggy-node-1"
# IP address that other nodes should use to connect to this node
ip = "127.0.0.1"
+# Explicit numeric replica ID for VSR consensus (0-based, must be unique per cluster)
+replica_id = 0
# Other nodes in the cluster
# Each field in 'ports' is optional. If omitted, the current node's configured port will be used.
@@ -559,7 +561,8 @@ ip = "127.0.0.1"
[[cluster.node.others]]
name = "iggy-node-2"
ip = "127.0.0.1"
-ports = { tcp = 8091, quic = 8081, http = 3001, websocket = 8093 }
+replica_id = 1
+ports = { tcp = 8091, quic = 8081, http = 3001, websocket = 8093, tcp_replica = 9091 }
# [[cluster.node.others]]
# name = "iggy-node-3"
diff --git a/core/shard/Cargo.toml b/core/shard/Cargo.toml
index 767c541c3..5d3151e7a 100644
--- a/core/shard/Cargo.toml
+++ b/core/shard/Cargo.toml
@@ -23,6 +23,7 @@ publish = false
[dependencies]
bytes = { workspace = true }
+compio = { workspace = true }
consensus = { path = "../consensus" }
crossfire = { workspace = true }
futures = { workspace = true }
diff --git a/core/shard/src/config.rs b/core/shard/src/config.rs
new file mode 100644
index 000000000..516582834
--- /dev/null
+++ b/core/shard/src/config.rs
@@ -0,0 +1,60 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Runtime tunables for the shard-0 coordinator.
+//!
+//! TODO: move this module into `core/configs` (as a `CoordinatorConfig`
+//! section nested under `ClusterConfig`) once downstream bootstrap
+//! wiring that constructs [`crate::coordinator::ShardZeroCoordinator`]
+//! from `ServerConfig` lands. Kept in-crate for now to avoid churning
+//! the configs crate ahead of that wiring.
+
+use std::time::Duration;
+
+/// Tunables for [`crate::coordinator::ShardZeroCoordinator`].
+#[derive(Debug, Clone)]
+pub struct CoordinatorConfig {
+ /// Cadence at which shard 0 re-broadcasts its authoritative
+ /// `replica_id -> owning_shard` snapshot. A shard whose inbox was
+ /// full when the original `ReplicaMappingUpdate` was sent recovers
+ /// its mapping on the next tick rather than staying silently stale.
+ /// Must be comfortably shorter than the VSR view-change timeout so
+ /// a missed mapping cannot trigger a spurious view change.
+ pub refresh_period: Duration,
+
+ /// When `total_shards > 1`, exclude shard 0 from replica placement.
+ /// Shard 0 already hosts the coordinator, the metadata writer, and
+ /// both listeners; replicas are long-lived steady flows, so offload
+ /// them to peer shards by default.
+ pub skip_shard_zero_for_replicas: bool,
+
+ /// When `total_shards > 1`, exclude shard 0 from client placement.
+ /// Default false: shard 0 continues to serve client traffic because
+ /// client connections are short-lived and benefit from shard-0
+ /// parallelism more than replicas do.
+ pub skip_shard_zero_for_clients: bool,
+}
+
+impl Default for CoordinatorConfig {
+ fn default() -> Self {
+ Self {
+ refresh_period: Duration::from_secs(10),
+ skip_shard_zero_for_replicas: true,
+ skip_shard_zero_for_clients: false,
+ }
+ }
+}
diff --git a/core/shard/src/coordinator.rs b/core/shard/src/coordinator.rs
new file mode 100644
index 000000000..5dc32c464
--- /dev/null
+++ b/core/shard/src/coordinator.rs
@@ -0,0 +1,561 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Shard-0 connection coordinator.
+//!
+//! Shard 0 is the sole binder of the replica listener and client listener,
+//! and the sole outbound dialer for higher-id peer replicas. On every
+//! accept / successful dial the coordinator:
+//!
+//! 1. picks the next target shard via round-robin,
+//! 2. duplicates the TCP fd,
+//! 3. sends a `ShardFramePayload::{Replica,Client}ConnectionSetup` frame to
+//! the target shard's inbox,
+//! 4. drops its own `TcpStream` so only the target shard's wrapped fd
+//! keeps the socket alive,
+//! 5. broadcasts the resulting replica -> owning shard mapping so every
+//! bus' `send_to_replica` slow path can route to the owner.
+//!
+//! Client ids encode the owning shard in their top 16 bits, so clients do
+//! not need a mapping table; any shard can route a client reply from the
+//! id alone.
+//!
+//! On `try_send` failure into the inter-shard channel (inbox full) the
+//! coordinator closes the duplicated fd and returns an error. VSR's
+//! retransmission plus the connector's periodic reconnect sweep cover the
+//! dropped connection.
+
+use crate::config::CoordinatorConfig;
+use crate::{ShardFrame, ShardFramePayload, TaggedSender, assert_sender_ordering};
+use compio::net::TcpStream;
+use compio::runtime::JoinHandle;
+use message_bus::SendError;
+use message_bus::fd_transfer;
+use message_bus::lifecycle::ShutdownToken;
+use std::cell::{Cell, RefCell};
+use std::rc::Rc;
+use std::time::Duration;
+use tracing::{debug, warn};
+
+/// Coordinator owned by shard 0 only.
+///
+/// Wrapped in `Rc` by the bootstrap and shared with the replica listener,
+/// the connector, and the client listener so each of those paths can
+/// delegate immediately.
+pub struct ShardZeroCoordinator<R: Send + 'static = ()> {
+ /// Inter-shard channel senders, indexed by shard id.
+ ///
+ /// Each [`TaggedSender`] carries the id of the shard whose paired
+ /// receiver drains it, and the ctor asserts `senders[i].shard_id() == i`.
+ /// The previous plain-`Sender` form was a silent-misroute hazard: a Vec
+ /// with correct length but permuted ordering had no way to be caught.
+ senders: Rc<Vec<TaggedSender<R>>>,
+ total_shards: u16,
+ cfg: CoordinatorConfig,
+ replica_rr: Cell<u16>,
+ client_rr: Cell<u16>,
+ client_seq: Cell<u128>,
+ /// Authoritative `replica_id -> owning_shard` view as understood by the
+ /// coordinator. Populated on successful `delegate_replica`, cleared by
+ /// [`forget_mapping`](Self::forget_mapping) when a replica dies. The
+ /// periodic refresh task re-broadcasts this snapshot so shards whose
+ /// inbox was full when the original `ReplicaMappingUpdate` was sent
+ /// recover their mapping on the next tick rather than staying silently
+ /// stale until the replica reconnects.
+ ///
+ /// Flat `[Option<u16>; 256]` to mirror `ReplicaRegistry` and avoid a
+ /// hash lookup on the refresh hot path. 512 bytes per coordinator.
+ mappings: RefCell<[Option<u16>; 256]>,
+}
+
+impl<R: Send + 'static> ShardZeroCoordinator<R> {
+ /// # Panics
+ ///
+ /// Panics if `senders.len() != total_shards`, if `total_shards < 1`, or
+ /// if any `senders[i].shard_id() != i`. Permuted senders are a
+ /// bootstrap programming error; `TaggedSender` lifts the ordering
+ /// invariant from a doc comment to a ctor assertion.
+ #[must_use]
+ pub fn new(
+ senders: Rc<Vec<TaggedSender<R>>>,
+ total_shards: u16,
+ cfg: CoordinatorConfig,
+ ) -> Self {
+ assert_eq!(
+ senders.len(),
+ total_shards as usize,
+ "senders must have one entry per shard",
+ );
+ assert!(total_shards >= 1, "total_shards must be at least 1");
+ assert_sender_ordering(&senders);
+ Self {
+ senders,
+ total_shards,
+ cfg,
+ replica_rr: Cell::new(0),
+ client_rr: Cell::new(0),
+ client_seq: Cell::new(1),
+ mappings: RefCell::new([None; 256]),
+ }
+ }
+
+ /// Pick the next target shard for a replica connection.
+ ///
+ /// When `total_shards > 1` and `cfg.skip_shard_zero_for_replicas` is
+ /// true (the default), the selection wraps over `[1, total_shards)`.
+ fn next_replica_target(&self) -> u16 {
+ rr_pick(
+ &self.replica_rr,
+ self.total_shards,
+ self.cfg.skip_shard_zero_for_replicas,
+ )
+ }
+
+ /// Pick the next target shard for a client connection.
+ ///
+ /// When `total_shards > 1` and `cfg.skip_shard_zero_for_clients` is
+ /// true, the selection wraps over `[1, total_shards)`. Default false.
+ fn next_client_target(&self) -> u16 {
+ rr_pick(
+ &self.client_rr,
+ self.total_shards,
+ self.cfg.skip_shard_zero_for_clients,
+ )
+ }
+
+ /// Mint a client id encoding `target_shard` in the top 16 bits and a
+ /// monotonic per-coordinator counter in the bottom 112 bits.
+ fn mint_client_id(&self, target_shard: u16) -> u128 {
+ let seq = self.client_seq.get();
+ self.client_seq.set(seq.wrapping_add(1));
+ (u128::from(target_shard) << 112) | seq
+ }
+
+ /// Ship a replica TCP connection to the next round-robin target shard.
+ ///
+ /// On success broadcasts a `ReplicaMappingUpdate` to every shard and
+ /// returns `Ok(target_shard)`. On inter-shard channel failure closes
+ /// the duplicated fd and returns `Err(SendError::RoutingFailed)`.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error when `dup(2)` fails or when the target shard's
+ /// inbox refuses the setup frame (full or disconnected).
+ pub fn delegate_replica(&self, stream: TcpStream, replica_id: u8) -> Result<u16, SendError> {
+ let target = self.next_replica_target();
+ let raw_fd = fd_transfer::dup_fd(&stream)
+ .map_err(|e| SendError::DupFailed(e.raw_os_error().unwrap_or(-1)))?;
+
+ let setup = ShardFramePayload::ReplicaConnectionSetup { raw_fd, replica_id };
+ if let Err(e) = self.senders[target as usize].try_send(ShardFrame::lifecycle(setup)) {
+ warn!(
+ replica_id,
+ target, "delegate_replica try_send failed: {e:?}"
+ );
+ fd_transfer::close_fd(raw_fd);
+ return Err(SendError::RoutingFailed(target));
+ }
+
+ // Shard 0 drops the original stream; the target shard's dup keeps
+ // the socket open.
+ drop(stream);
+
+ // Record the mapping on the coordinator's authoritative snapshot
+ // BEFORE broadcasting: if `broadcast_mapping_update` drops a frame
+ // on a full inbox, the periodic refresh task reads this snapshot
+ // and re-broadcasts.
+ self.mappings.borrow_mut()[usize::from(replica_id)] = Some(target);
+ self.broadcast_mapping_update(replica_id, target);
+
+ Ok(target)
+ }
+
+ fn broadcast_mapping_update(&self, replica_id: u8, owning_shard: u16) {
+ // Broadcast to every shard (including shard 0 and the owner
+ // itself). Drops on full inbox are tolerable: the owning shard
+ // already holds the fd and the periodic refresh task reconciles
+ // any missed mapping on its next tick.
+ for sender in self.senders.iter() {
+ let update = ShardFramePayload::ReplicaMappingUpdate {
+ replica_id,
+ owning_shard,
+ };
+ if let Err(e) = sender.try_send(ShardFrame::lifecycle(update)) {
+ debug!(
+ shard_id = sender.shard_id(),
+ replica_id, "mapping update try_send failed: {e:?}"
+ );
+ }
+ }
+ }
+
+ /// Ship a client TCP connection to the next round-robin target shard.
+ ///
+ /// On success returns the minted client id. On failure closes the
+ /// duplicated fd and returns an error.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error when `dup(2)` fails or the target shard's inbox
+ /// refuses the setup frame.
+ 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 raw_fd = fd_transfer::dup_fd(&stream)
+ .map_err(|e| SendError::DupFailed(e.raw_os_error().unwrap_or(-1)))?;
+ let setup = ShardFramePayload::ClientConnectionSetup { raw_fd, client_id };
+ if let Err(e) = self.senders[target as usize].try_send(ShardFrame::lifecycle(setup)) {
+ warn!(client_id, target, "delegate_client try_send failed: {e:?}");
+ fd_transfer::close_fd(raw_fd);
+ return Err(SendError::RoutingFailed(target));
+ }
+
+ drop(stream);
+ Ok(client_id)
+ }
+
+ /// Broadcast a `ReplicaMappingClear` to every shard. Used by the
+ /// `ConnectionLost` handler before the next `delegate_replica` runs.
+ pub fn broadcast_mapping_clear(&self, replica_id: u8) {
+ crate::broadcast_mapping_clear(&self.senders, replica_id);
+ }
+
+ /// Drop the coordinator's authoritative entry for `replica_id`, so
+ /// the periodic refresh task does not keep re-broadcasting the
+ /// mapping of a replica that has disconnected.
+ ///
+ /// Call from shard 0's `ConnectionLost` handler paired with
+ /// [`broadcast_mapping_clear`](Self::broadcast_mapping_clear).
+ pub fn forget_mapping(&self, replica_id: u8) {
+ self.mappings.borrow_mut()[usize::from(replica_id)] = None;
+ }
+
+ /// Re-broadcast every mapping the coordinator is tracking.
+ ///
+ /// A `delegate_replica` broadcast is `try_send` best-effort: a shard
+ /// whose inbox was full at the moment of the original broadcast
+ /// permanently lost that mapping and `send_to_replica` on that shard
+ /// would return `ReplicaNotConnected` until the next reconnect sweep
+ /// reinstalled the peer. Periodic refresh covers that gap.
+ pub fn broadcast_mapping_snapshot(&self) {
+ // Collect into a small stack buffer to release the `RefCell`
+ // borrow before broadcasting (try_send is sync but the borrow is
+ // cheap to drop first for consistency with other mapping paths).
+ let snapshot: Vec<(u8, u16)> = self
+ .mappings
+ .borrow()
+ .iter()
+ .enumerate()
+ .filter_map(|(idx, slot)| {
+ slot.map(|owner| {
+ #[allow(clippy::cast_possible_truncation)]
+ (idx as u8, owner)
+ })
+ })
+ .collect();
+ for (replica_id, owning_shard) in snapshot {
+ self.broadcast_mapping_update(replica_id, owning_shard);
+ }
+ }
+
+ /// Spawn a compio task that calls [`broadcast_mapping_snapshot`] every
+ /// `period` until `token` fires. Bootstrap is expected to track the
+ /// returned handle on the bus's background tasks so graceful shutdown
+ /// awaits it.
+ pub fn spawn_refresh_task(
+ self: &Rc<Self>,
+ token: ShutdownToken,
+ period: Duration,
+ ) -> JoinHandle<()> {
+ let coord = Rc::clone(self);
+ compio::runtime::spawn(async move {
+ while token.sleep_or_shutdown(period).await {
+ coord.broadcast_mapping_snapshot();
+ }
+ debug!("coordinator mapping refresh task exiting");
+ })
+ }
+
+ #[must_use]
+ pub const fn total_shards(&self) -> u16 {
+ self.total_shards
+ }
+}
+
+/// Advance `counter` and return the next target shard.
+///
+/// When `skip_zero` is true and `total_shards > 1`, wraps over
+/// `[1, total_shards)`; otherwise wraps over `[0, total_shards)`. With
+/// `total_shards == 1` the flag is ignored and the function always
+/// returns 0.
+fn rr_pick(counter: &Cell<u16>, total_shards: u16, skip_zero: bool) -> u16 {
+ let use_skip = skip_zero && total_shards > 1;
+ let offset: u16 = u16::from(use_skip);
+ let width = total_shards.saturating_sub(offset).max(1);
+ let cur = counter.get();
+ counter.set(cur.wrapping_add(1));
+ offset + (cur % width)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use compio::net::{TcpListener, TcpStream};
+
+ fn build_senders(total: u16) -> Rc<Vec<TaggedSender>> {
+ let mut senders = Vec::with_capacity(total as usize);
+ for shard_id in 0..total {
+ let (tx, _rx) = crate::shard_channel::<()>(shard_id, 16);
+ senders.push(tx);
+ }
+ Rc::new(senders)
+ }
+
+ fn build_senders_with_rx(
+ total: u16,
+ ) -> (Rc<Vec<TaggedSender>>, Vec<crate::Receiver<ShardFrame>>) {
+ let mut senders = Vec::with_capacity(total as usize);
+ let mut receivers = Vec::with_capacity(total as usize);
+ for shard_id in 0..total {
+ let (tx, rx) = crate::shard_channel::<()>(shard_id, 16);
+ senders.push(tx);
+ receivers.push(rx);
+ }
+ (Rc::new(senders), receivers)
+ }
+
+ #[test]
+ #[should_panic(expected = "inter-shard vec must be in canonical order")]
+ fn ctor_rejects_permuted_sender_vec() {
+ // Build senders in correct order, then swap two entries so the
+ // indexed position no longer matches the tagged shard id.
+ let (tx0, _rx0) = crate::shard_channel::<()>(0, 16);
+ let (tx1, _rx1) = crate::shard_channel::<()>(1, 16);
+ let (tx2, _rx2) = crate::shard_channel::<()>(2, 16);
+ let (tx3, _rx3) = crate::shard_channel::<()>(3, 16);
+ let permuted = Rc::new(vec![tx0, tx2, tx1, tx3]);
+ let _coord = ShardZeroCoordinator::<()>::new(permuted, 4, CoordinatorConfig::default());
+ }
+
+ #[test]
+ fn replica_rr_default_skips_shard_zero() {
+ // Default config: replicas skip shard 0, clients include it.
+ let senders = build_senders(4);
+ let coord = ShardZeroCoordinator::<()>::new(senders, 4, CoordinatorConfig::default());
+
+ // Replicas wrap over [1, 4).
+ assert_eq!(coord.next_replica_target(), 1);
+ assert_eq!(coord.next_replica_target(), 2);
+ assert_eq!(coord.next_replica_target(), 3);
+ assert_eq!(coord.next_replica_target(), 1);
+ assert_eq!(coord.next_replica_target(), 2);
+
+ // Clients span [0, 4).
+ assert_eq!(coord.next_client_target(), 0);
+ assert_eq!(coord.next_client_target(), 1);
+ assert_eq!(coord.next_client_target(), 2);
+ assert_eq!(coord.next_client_target(), 3);
+ assert_eq!(coord.next_client_target(), 0);
+ }
+
+ #[test]
+ fn rr_includes_shard_zero_when_skip_flags_off() {
+ let senders = build_senders(4);
+ let cfg = CoordinatorConfig {
+ skip_shard_zero_for_replicas: false,
+ skip_shard_zero_for_clients: false,
+ ..CoordinatorConfig::default()
+ };
+ let coord = ShardZeroCoordinator::<()>::new(senders, 4, cfg);
+
+ assert_eq!(coord.next_replica_target(), 0);
+ assert_eq!(coord.next_replica_target(), 1);
+ assert_eq!(coord.next_replica_target(), 2);
+ assert_eq!(coord.next_replica_target(), 3);
+ assert_eq!(coord.next_replica_target(), 0);
+
+ assert_eq!(coord.next_client_target(), 0);
+ assert_eq!(coord.next_client_target(), 1);
+ }
+
+ #[test]
+ fn rr_skips_shard_zero_when_both_flags_on() {
+ let senders = build_senders(4);
+ let cfg = CoordinatorConfig {
+ skip_shard_zero_for_replicas: true,
+ skip_shard_zero_for_clients: true,
+ ..CoordinatorConfig::default()
+ };
+ let coord = ShardZeroCoordinator::<()>::new(senders, 4, cfg);
+
+ for _ in 0..8 {
+ let r = coord.next_replica_target();
+ let c = coord.next_client_target();
+ assert!((1..4).contains(&r), "replica target {r} must skip shard 0");
+ assert!((1..4).contains(&c), "client target {c} must skip shard 0");
+ }
+ }
+
+ #[test]
+ fn rr_single_shard_returns_zero_regardless_of_flags() {
+ let senders = build_senders(1);
+ let cfg = CoordinatorConfig {
+ skip_shard_zero_for_replicas: true,
+ skip_shard_zero_for_clients: true,
+ ..CoordinatorConfig::default()
+ };
+ let coord = ShardZeroCoordinator::<()>::new(senders, 1, cfg);
+ for _ in 0..4 {
+ assert_eq!(coord.next_replica_target(), 0);
+ assert_eq!(coord.next_client_target(), 0);
+ }
+ }
+
+ #[test]
+ fn mint_client_id_encodes_target_shard() {
+ let senders = build_senders(8);
+ let coord = ShardZeroCoordinator::<()>::new(senders, 8, CoordinatorConfig::default());
+
+ let id = coord.mint_client_id(5);
+ assert_eq!((id >> 112) as u16, 5);
+ assert_eq!(id & ((1u128 << 112) - 1), 1, "first seq is 1");
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn delegate_replica_sends_setup_then_broadcasts_mapping() {
+ let (senders, receivers) = build_senders_with_rx(4);
+ let coord = ShardZeroCoordinator::<()>::new(senders, 4, CoordinatorConfig::default());
+
+ // Loopback TCP pair so delegate_replica has a real fd to dup.
+ let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = listener.local_addr().unwrap();
+ let accept = compio::runtime::spawn(async move { listener.accept().await.unwrap() });
+ let client = TcpStream::connect(addr).await.unwrap();
+ let (_server, _peer_addr) = accept.await.unwrap();
+
+ let target = coord.delegate_replica(client, 7).expect("delegate ok");
+ assert_eq!(
+ target, 1,
+ "first replica target skips shard 0 under default config",
+ );
+
+ // Target shard should observe ReplicaConnectionSetup.
+ let setup_frame = receivers[target as usize].recv().await.unwrap();
+ match setup_frame.payload {
+ ShardFramePayload::ReplicaConnectionSetup { raw_fd, replica_id } => {
+ assert_eq!(replica_id, 7);
+ // Clean up the dup'd fd so the test does not leak.
+ fd_transfer::close_fd(raw_fd);
+ }
+ _ => panic!("expected ReplicaConnectionSetup on target shard"),
+ }
+
+ // Every shard (including shard 0 and the target itself) should
+ // observe the ReplicaMappingUpdate broadcast.
+ for (idx, rx) in receivers.iter().enumerate() {
+ let frame = rx.recv().await.unwrap();
+ match frame.payload {
+ ShardFramePayload::ReplicaMappingUpdate {
+ replica_id,
+ owning_shard,
+ } => {
+ assert_eq!(replica_id, 7, "shard {idx} mapping update replica id");
+ assert_eq!(owning_shard, target, "shard {idx} mapping update owner");
+ }
+ _ => panic!("shard {idx} expected ReplicaMappingUpdate"),
+ }
+ }
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn broadcast_mapping_clear_reaches_every_shard() {
+ let (senders, receivers) = build_senders_with_rx(4);
+ let coord = ShardZeroCoordinator::<()>::new(senders, 4, CoordinatorConfig::default());
+
+ coord.broadcast_mapping_clear(9);
+
+ for (idx, rx) in receivers.iter().enumerate() {
+ let frame = rx.recv().await.unwrap();
+ match frame.payload {
+ ShardFramePayload::ReplicaMappingClear { replica_id } => {
+ assert_eq!(replica_id, 9, "shard {idx} mapping clear replica id");
+ }
+ _ => panic!("shard {idx} expected ReplicaMappingClear"),
+ }
+ }
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn snapshot_rebroadcasts_every_tracked_mapping() {
+ let (senders, receivers) = build_senders_with_rx(4);
+ let coord = Rc::new(ShardZeroCoordinator::<()>::new(
+ senders,
+ 4,
+ CoordinatorConfig::default(),
+ ));
+
+ // Seed the coordinator's tracked mappings directly (delegate_replica
+ // needs a real TCP fd; the snapshot path is orthogonal to dup).
+ coord.mappings.borrow_mut()[3] = Some(1);
+ coord.mappings.borrow_mut()[7] = Some(2);
+
+ coord.broadcast_mapping_snapshot();
+
+ // Each shard must observe an update for each seeded mapping. The
+ // order across replica ids is unspecified (HashMap->Vec iteration
+ // on the collected snapshot), so collect and compare as a set.
+ for (idx, rx) in receivers.iter().enumerate() {
+ let mut observed = std::collections::BTreeSet::new();
+ for _ in 0..2 {
+ let frame = rx.recv().await.unwrap();
+ match frame.payload {
+ ShardFramePayload::ReplicaMappingUpdate {
+ replica_id,
+ owning_shard,
+ } => {
+ observed.insert((replica_id, owning_shard));
+ }
+ _ => panic!("shard {idx} expected ReplicaMappingUpdate"),
+ }
+ }
+ let expected: std::collections::BTreeSet<_> =
+ [(3u8, 1u16), (7u8, 2u16)].into_iter().collect();
+ assert_eq!(
+ observed, expected,
+ "shard {idx} did not receive the full snapshot"
+ );
+ }
+ }
+
+ #[compio::test]
+ #[allow(clippy::future_not_send)]
+ async fn forget_mapping_prunes_entry_from_future_snapshots() {
+ let (senders, _receivers) = build_senders_with_rx(2);
+ let coord = ShardZeroCoordinator::<()>::new(senders, 2, CoordinatorConfig::default());
+
+ coord.mappings.borrow_mut()[4] = Some(1);
+ coord.mappings.borrow_mut()[5] = Some(1);
+ coord.forget_mapping(4);
+
+ assert!(coord.mappings.borrow()[4].is_none());
+ assert_eq!(coord.mappings.borrow()[5], Some(1));
+ }
+}
diff --git a/core/shard/src/lib.rs b/core/shard/src/lib.rs
index 47bd5a9b8..648423645 100644
--- a/core/shard/src/lib.rs
+++ b/core/shard/src/lib.rs
@@ -15,13 +15,19 @@
// specific language governing permissions and limitations
// under the License.
+pub mod config;
+pub mod coordinator;
mod router;
pub mod shards_table;
+pub use config::CoordinatorConfig;
+
use consensus::{
LocalPipeline, MetadataHandle, MuxPlane, PartitionsHandle, Pipeline, Plane, PlaneKind,
Sequencer, VsrAction, VsrConsensus,
};
+use iggy_binary_protocol::consensus::MESSAGE_ALIGN;
+use iggy_binary_protocol::consensus::iobuf::Frozen;
use iggy_binary_protocol::{
Command2, CommitHeader, DoViewChangeHeader, GenericHeader, Message, MessageBag, PrepareHeader,
PrepareOkHeader, RequestHeader, StartViewChangeHeader, StartViewHeader,
@@ -30,10 +36,13 @@ use iggy_common::variadic;
use iggy_common::{PartitionStats, sharding::IggyNamespace};
use journal::{Journal, JournalHandle};
use message_bus::MessageBus;
+use message_bus::client_listener::RequestHandler;
+use message_bus::replica_listener::MessageHandler;
use metadata::IggyMetadata;
use metadata::stm::StateMachine;
use partitions::{IggyPartition, IggyPartitions};
use shards_table::ShardsTable;
+use std::rc::Rc;
use std::sync::Arc;
pub type ShardPlane<B, J, S, M> =
@@ -86,9 +95,122 @@ pub fn channel<T: Send + 'static>(capacity: usize) -> (Sender<T>, Receiver<T>) {
crossfire::mpsc::bounded_blocking_async(capacity)
}
+/// Create a bounded inter-shard channel whose sender is tagged with the
+/// owning shard.
+///
+/// Bootstrap uses this to build the per-shard sender `Vec` such that
+/// `vec[i]` necessarily reaches shard `i`.
+#[must_use]
+pub fn shard_channel<R: Send + 'static>(
+ owner_shard: u16,
+ capacity: usize,
+) -> (TaggedSender<R>, Receiver<ShardFrame<R>>) {
+ let (tx, rx) = channel::<ShardFrame<R>>(capacity);
+ (TaggedSender::new(owner_shard, tx), rx)
+}
+
+/// A [`Sender`] annotated with the id of the shard whose paired receiver it
+/// feeds.
+///
+/// Inter-shard routing indexes `senders[i]` with `i == target_shard`. The
+/// plain `Sender` form has no way to verify that invariant at runtime, so a
+/// permuted `Vec<Sender<_>>` would silently misroute every setup, mapping,
+/// and forward frame. Construct senders through [`shard_channel`] (or
+/// [`TaggedSender::new`]) at the channel-creation site; the coordinator and
+/// [`IggyShard`] ctors then assert `senders[i].shard_id() == i`, turning the
+/// ordering invariant from a doc comment into a ctor-checked type property.
+pub struct TaggedSender<R: Send + 'static = ()> {
+ shard_id: u16,
+ inner: Sender<ShardFrame<R>>,
+}
+
+impl<R: Send + 'static> TaggedSender<R> {
+ /// Wrap an already-constructed sender with the id of the shard whose
+ /// paired receiver drains it. Prefer [`shard_channel`] unless an
+ /// existing sender is being re-tagged (e.g., tests that build senders
+ /// manually and know the ordering is correct).
+ #[must_use]
+ pub const fn new(shard_id: u16, inner: Sender<ShardFrame<R>>) -> Self {
+ Self { shard_id, inner }
+ }
+
+ #[must_use]
+ pub const fn shard_id(&self) -> u16 {
+ self.shard_id
+ }
+
+ #[must_use]
+ pub const fn sender(&self) -> &Sender<ShardFrame<R>> {
+ &self.inner
+ }
+}
+
+impl<R: Send + 'static> std::ops::Deref for TaggedSender<R> {
+ type Target = Sender<ShardFrame<R>>;
+
+ fn deref(&self) -> &Self::Target {
+ &self.inner
+ }
+}
+
+/// Assert the canonical ordering `senders[i].shard_id() == i`. Called from
+/// every ctor that accepts a pre-built `Vec<TaggedSender<_>>`.
+///
+/// # Panics
+///
+/// Panics if any sender in `senders` carries a `shard_id` that does not
+/// match its index. Bootstrap programming error.
+fn assert_sender_ordering<R: Send + 'static>(senders: &[TaggedSender<R>]) {
+ for (idx, sender) in senders.iter().enumerate() {
+ let expected = u16::try_from(idx).expect("shard count must fit in u16");
+ assert_eq!(
+ sender.shard_id(),
+ expected,
+ "senders[{idx}] carries shard_id {}; inter-shard vec must be in canonical order",
+ sender.shard_id(),
+ );
+ }
+}
+
+/// Payload carried by a [`ShardFrame`].
+pub enum ShardFramePayload {
+ /// A consensus protocol message routed between shards.
+ Consensus(Message<GenericHeader>),
+ /// Shard 0 distributes an inbound replica TCP connection fd to the owning
+ /// shard. The receiving shard wraps the fd in a `TcpStream` and spawns
+ /// writer + reader tasks on its own compio runtime.
+ ReplicaConnectionSetup { raw_fd: i32, replica_id: u8 },
+ /// Shard 0 distributes an inbound SDK client TCP connection fd to the
+ /// owning shard. The receiving shard wraps the fd and installs client
+ /// reader / writer tasks locally. The owning shard is encoded in the top
+ /// 16 bits of `client_id`.
+ ClientConnectionSetup { raw_fd: i32, client_id: u128 },
+ /// Shard 0 broadcasts the owner for a replica to every shard so each
+ /// bus' `send_to_replica` slow path can route through the correct owner.
+ ReplicaMappingUpdate { replica_id: u8, owning_shard: u16 },
+ /// Shard 0 broadcasts that a replica mapping should be forgotten (e.g.
+ /// after a connection loss and before the next allocate).
+ ReplicaMappingClear { replica_id: u8 },
+ /// A non-owning shard forwards a replica send to the owning shard's
+ /// local bus; the owning shard then takes the fast path.
+ ForwardReplicaSend {
+ replica_id: u8,
+ msg: Frozen<MESSAGE_ALIGN>,
+ },
+ /// A shard that doesn't hold the client's TCP connection forwards a
+ /// client send to the owning shard (top 16 bits of `client_id`).
+ ForwardClientSend {
+ client_id: u128,
+ msg: Frozen<MESSAGE_ALIGN>,
+ },
+ /// Owning shard notifies shard 0 that a replica connection died.
+ /// Shard 0 clears the mapping cluster-wide and drives a reconnect.
+ ConnectionLost { replica_id: u8 },
+}
+
/// Envelope for inter-shard channel messages.
///
-/// Wraps a consensus [`Message`] together with an optional one-shot response
+/// Wraps a [`ShardFramePayload`] together with an optional one-shot response
/// channel. Fire-and-forget dispatches leave `response_sender` as `None`;
/// request-response dispatches provide a sender that the message pump will
/// notify once the message has been processed.
@@ -96,33 +218,65 @@ pub fn channel<T: Send + 'static>(capacity: usize) -> (Sender<T>, Receiver<T>) {
/// The response type `R` is generic so that higher layers (e.g. HTTP handlers)
/// can carry a response enum while the consensus layer can default to `()`.
pub struct ShardFrame<R: Send + 'static = ()> {
- pub message: Message<GenericHeader>,
+ pub payload: ShardFramePayload,
pub response_sender: Option<Sender<R>>,
}
impl<R: Send + 'static> ShardFrame<R> {
- /// Create a fire-and-forget frame (no caller waiting for completion).
+ /// Create a fire-and-forget consensus message frame.
#[must_use]
pub const fn fire_and_forget(message: Message<GenericHeader>) -> Self {
Self {
- message,
+ payload: ShardFramePayload::Consensus(message),
response_sender: None,
}
}
- /// Create a request-response frame. Returns the frame and a receiver
- /// that the caller can await for completion notification.
+ /// Create a request-response consensus message frame. Returns the frame
+ /// and a receiver that the caller can await for completion notification.
#[must_use]
pub fn with_response(message: Message<GenericHeader>) -> (Self, Receiver<R>) {
let (tx, rx) = channel(1);
(
Self {
- message,
+ payload: ShardFramePayload::Consensus(message),
response_sender: Some(tx),
},
rx,
)
}
+
+ /// Create a fire-and-forget lifecycle frame (connection setup, loss).
+ #[must_use]
+ pub const fn lifecycle(payload: ShardFramePayload) -> Self {
+ Self {
+ payload,
+ response_sender: None,
+ }
+ }
+}
+
+/// Broadcast a `ReplicaMappingClear` to every shard (including sender-self).
+///
+/// Used by shard 0's `ConnectionLost` handler and by
+/// [`coordinator::ShardZeroCoordinator::broadcast_mapping_clear`] so both
+/// paths go through the same try-send-and-log logic. `try_send` failures
+/// (inbox full or disconnected) are logged at debug: the next mapping
+/// broadcast or reconnect sweep will reconcile.
+pub(crate) fn broadcast_mapping_clear<R: Send + 'static>(
+ senders: &[TaggedSender<R>],
+ replica_id: u8,
+) {
+ for sender in senders {
+ let clear = ShardFramePayload::ReplicaMappingClear { replica_id };
+ if let Err(e) = sender.try_send(ShardFrame::lifecycle(clear)) {
+ tracing::debug!(
+ shard_id = sender.shard_id(),
+ replica_id,
+ "mapping clear try_send failed: {e:?}"
+ );
+ }
+ }
}
pub struct IggyShard<B, MJ, S, M, T = (), R: Send + 'static = ()>
@@ -133,10 +287,28 @@ where
pub name: String,
pub plane: ShardPlane<B, MJ, S, M>,
+ /// Handle to the local bus. Retained alongside the bus owned by every
+ /// consensus plane so the router can reach the `ConnectionInstaller` /
+ /// mapping surface without going through consensus.
+ pub bus: B,
+
+ /// Callback attached to every delegated replica connection installed
+ /// on this shard. The bus' reader task invokes this for each inbound
+ /// consensus message; the callback is typically `|_, msg| shard.dispatch(msg)`.
+ on_replica_message: MessageHandler,
+
+ /// Callback attached to every delegated client connection installed on
+ /// this shard. Invoked for each inbound `Request` frame.
+ on_client_request: RequestHandler,
+
/// Channel senders to every shard, indexed by shard id.
/// Includes a sender to self so that local routing goes through the
/// same channel path as remote routing.
- senders: Vec<Sender<ShardFrame<R>>>,
+ ///
+ /// [`assert_sender_ordering`] is invoked in the ctor so `senders[i]`
+ /// is guaranteed to feed the shard whose `id == i`. Call sites can
+ /// therefore index by `target_shard` without re-checking.
+ senders: Vec<TaggedSender<R>>,
/// Receiver end of this shard's inbox. Peer shards (and self) send
/// messages here via the corresponding sender.
@@ -146,6 +318,13 @@ where
shards_table: T,
partition_consensus: PartitionConsensusConfig<B>,
+
+ /// Shard 0 coordinator, set at bootstrap on shard 0 only. The router's
+ /// `ConnectionLost` handler calls [`coordinator::ShardZeroCoordinator::forget_mapping`]
+ /// so the periodic refresh task stops re-broadcasting the mapping of
+ /// a replica that has disconnected. `None` on non-zero shards and in
+ /// single-shard tests that bypass the coordinator.
+ coordinator: Option<Rc<crate::coordinator::ShardZeroCoordinator<R>>>,
}
impl<B, MJ, S, M, T, R: Send + 'static> IggyShard<B, MJ, S, M, T, R>
@@ -155,39 +334,74 @@ where
{
/// Create a new shard with channel links and a shards table.
///
- /// * `senders` - one sender per shard in the cluster (indexed by shard id).
+ /// * `bus` - shard-local bus handle (kept alongside the buses owned
+ /// by the consensus planes so the router can reach installer /
+ /// mapping operations directly).
+ /// * `senders` - one [`TaggedSender`] per shard. The ctor asserts
+ /// `senders[i].shard_id() == i`; use [`shard_channel`] at
+ /// construction time so every sender carries the id of the shard
+ /// whose receiver drains it.
/// * `inbox` - the receiver that this shard drains in its message pump.
/// * `shards_table` - namespace -> shard routing table.
+ ///
+ /// # Panics
+ ///
+ /// Panics if `senders` is not in canonical order (any
+ /// `senders[i].shard_id() != i`). That is a bootstrap programming
+ /// error; the resulting permutation would silently misroute every
+ /// inter-shard frame.
#[must_use]
+ #[allow(clippy::too_many_arguments)]
pub fn new(
identity: ShardIdentity,
+ bus: B,
+ on_replica_message: MessageHandler,
+ on_client_request: RequestHandler,
metadata: IggyMetadata<VsrConsensus<B>, MJ, S, M>,
partitions: IggyPartitions<B>,
- senders: Vec<Sender<ShardFrame<R>>>,
+ senders: Vec<TaggedSender<R>>,
inbox: Receiver<ShardFrame<R>>,
shards_table: T,
partition_consensus: PartitionConsensusConfig<B>,
) -> Self {
+ assert_sender_ordering(&senders);
let plane = MuxPlane::new(variadic!(metadata, partitions));
let ShardIdentity { id, name } = identity;
Self {
id,
name,
plane,
+ bus,
+ on_replica_message,
+ on_client_request,
senders,
inbox,
shards_table,
partition_consensus,
+ coordinator: None,
}
}
- /// Create a shard without inter-shard channels.
+ /// Attach a shard-0 coordinator. Bootstrap calls this on the shard 0
+ /// `IggyShard` only; other shards keep `coordinator = None`. The
+ /// router's `ConnectionLost` handler will call
+ /// [`coordinator::ShardZeroCoordinator::forget_mapping`] via this
+ /// reference so the periodic refresh stops re-broadcasting dead
+ /// replica mappings.
+ pub fn set_coordinator(&mut self, coord: Rc<crate::coordinator::ShardZeroCoordinator<R>>) {
+ self.coordinator = Some(coord);
+ }
+
+ /// Create a shard without inter-shard channels or delegated connections.
///
- /// Useful for the simulator where inbound messages are delivered directly
- /// via [`on_message`](Self::on_message) instead of through an inbox channel.
+ /// Useful for the simulator where inbound messages are delivered
+ /// directly via [`on_message`](Self::on_message) instead of the TCP /
+ /// fd-transfer path. Installs no-op connection handlers because the
+ /// simulator never receives a `ReplicaConnectionSetup` frame.
#[must_use]
pub fn without_inbox(
identity: ShardIdentity,
+ bus: B,
metadata: IggyMetadata<VsrConsensus<B>, MJ, S, M>,
partitions: IggyPartitions<B>,
shards_table: T,
@@ -202,7 +416,11 @@ where
Self {
id,
name,
+ bus,
+ on_replica_message: std::rc::Rc::new(|_, _| {}),
+ on_client_request: std::rc::Rc::new(|_, _| {}),
plane,
+ coordinator: None,
senders: Vec::new(),
inbox,
shards_table,
@@ -229,7 +447,7 @@ where
#[allow(clippy::future_not_send)]
pub async fn on_message(&self, message: Message<GenericHeader>)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -259,7 +477,7 @@ where
#[allow(clippy::future_not_send)]
pub async fn on_request(&self, request: Message<RequestHeader>)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -278,7 +496,7 @@ where
#[allow(clippy::future_not_send)]
pub async fn on_replicate(&self, prepare: Message<PrepareHeader>)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -297,7 +515,7 @@ where
#[allow(clippy::future_not_send)]
pub async fn on_ack(&self, prepare_ok: Message<PrepareOkHeader>)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -327,7 +545,7 @@ where
#[allow(clippy::future_not_send)]
pub async fn process_loopback(&self, buf: &mut Vec<Message<GenericHeader>>) -> usize
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -385,11 +603,7 @@ where
/// by the partition consensus replica id.
pub fn init_partition(&mut self, namespace: IggyNamespace)
where
- B: MessageBus<
- Replica = u8,
- Data = iggy_binary_protocol::Message<iggy_binary_protocol::GenericHeader>,
- Client = u128,
- > + Clone,
+ B: MessageBus + Clone,
{
let partitions = self.plane.partitions_mut();
if partitions.contains(&namespace) {
@@ -423,7 +637,7 @@ where
#[allow(clippy::future_not_send)]
async fn on_start_view_change(&self, msg: Message<StartViewChangeHeader>)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -470,7 +684,7 @@ where
#[allow(clippy::future_not_send)]
async fn on_do_view_change(&self, msg: Message<DoViewChangeHeader>)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -535,7 +749,7 @@ where
#[allow(clippy::future_not_send)]
async fn on_start_view(&self, msg: Message<StartViewHeader>)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -582,7 +796,7 @@ where
#[allow(clippy::future_not_send)]
async fn on_commit(&self, msg: &Message<CommitHeader>)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -637,7 +851,7 @@ where
#[allow(clippy::future_not_send)]
pub async fn tick_partitions(&self)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -665,7 +879,7 @@ where
#[allow(clippy::future_not_send)]
pub async fn tick_metadata(&self)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -698,7 +912,7 @@ async fn dispatch_vsr_actions<B, P, J>(
journal: Option<&J>,
actions: &[VsrAction],
) where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = consensus::PipelineEntry>,
J: JournalHandle,
<J as JournalHandle>::Target: Journal<
@@ -714,12 +928,22 @@ async fn dispatch_vsr_actions<B, P, J>(
let cluster = consensus.cluster();
let replica_count = consensus.replica_count();
- let send = |target: u8, msg: Message<GenericHeader>| async move {
+ let send = |target: u8, msg: Frozen<MESSAGE_ALIGN>| async move {
if let Err(e) = bus.send_to_replica(target, msg).await {
tracing::debug!(replica = self_id, target, "bus send failed: {e}");
}
};
+ let broadcast = async |frozen: Frozen<MESSAGE_ALIGN>| {
+ // Freeze once at the primary; each target just bumps the atomic
+ // refcount on the underlying ControlBlock.
+ for target in 0..replica_count {
+ if target != self_id {
+ send(target, frozen.clone()).await;
+ }
+ }
+ };
+
for action in actions {
match action {
VsrAction::SendStartViewChange { view, namespace } => {
@@ -732,11 +956,7 @@ async fn dispatch_vsr_actions<B, P, J>(
h.namespace = *namespace;
h.size = size_of::<StartViewChangeHeader>() as u32;
});
- for target in 0..replica_count {
- if target != self_id {
- send(target, msg.deep_copy().into_generic()).await;
- }
- }
+ broadcast(msg.into_generic().into_frozen()).await;
}
VsrAction::SendDoViewChange {
view,
@@ -758,7 +978,7 @@ async fn dispatch_vsr_actions<B, P, J>(
h.namespace = *namespace;
h.size = size_of::<DoViewChangeHeader>() as u32;
});
- send(*target, msg.into_generic()).await;
+ send(*target, msg.into_generic().into_frozen()).await;
}
VsrAction::SendStartView {
view,
@@ -777,11 +997,7 @@ async fn dispatch_vsr_actions<B, P, J>(
h.namespace = *namespace;
h.size = size_of::<StartViewHeader>() as u32;
});
- for target in 0..replica_count {
- if target != self_id {
- send(target, msg.deep_copy().into_generic()).await;
- }
- }
+ broadcast(msg.into_generic().into_frozen()).await;
}
VsrAction::SendPrepareOk {
view,
@@ -814,7 +1030,7 @@ async fn dispatch_vsr_actions<B, P, J>(
h.namespace = *namespace;
h.size = size_of::<PrepareOkHeader>() as u32;
});
- send(*target, msg.into_generic()).await;
+ send(*target, msg.into_generic().into_frozen()).await;
}
}
VsrAction::RetransmitPrepares { targets } => {
@@ -825,8 +1041,10 @@ async fn dispatch_vsr_actions<B, P, J>(
let Some(prepare) = journal.handle().entry(header).await else {
continue;
};
+ // Freeze the retransmit payload once; clone per target.
+ let frozen = prepare.into_generic().into_frozen();
for replica in replicas {
- send(*replica, prepare.clone().into_generic()).await;
+ send(*replica, frozen.clone()).await;
}
}
}
@@ -894,11 +1112,7 @@ async fn dispatch_vsr_actions<B, P, J>(
h.size = size_of::<CommitHeader>() as u32;
},
);
- for target in 0..replica_count {
- if target != self_id {
- send(target, msg.deep_copy().into_generic()).await;
- }
- }
+ broadcast(msg.into_generic().into_frozen()).await;
}
}
}
@@ -914,7 +1128,7 @@ async fn dispatch_partition_journal_actions<B, P>(
partition: &IggyPartition<B>,
actions: &[VsrAction],
) where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
P: Pipeline<Entry = consensus::PipelineEntry>,
{
use std::mem::size_of;
@@ -924,7 +1138,7 @@ async fn dispatch_partition_journal_actions<B, P>(
let cluster = consensus.cluster();
let journal = &partition.log.journal().inner;
- let send = |target: u8, msg: Message<GenericHeader>| async move {
+ let send = |target: u8, msg: Frozen<MESSAGE_ALIGN>| async move {
if let Err(e) = bus.send_to_replica(target, msg).await {
tracing::debug!(replica = self_id, target, "bus send failed: {e}");
}
@@ -959,7 +1173,7 @@ async fn dispatch_partition_journal_actions<B, P>(
h.namespace = *namespace;
h.size = size_of::<PrepareOkHeader>() as u32;
});
- send(*target, msg.into_generic()).await;
+ send(*target, msg.into_generic().into_frozen()).await;
}
}
VsrAction::RetransmitPrepares { targets } => {
@@ -967,14 +1181,15 @@ async fn dispatch_partition_journal_actions<B, P>(
let Some(prepare) = journal.entry(header).await else {
continue;
};
- let prepare = Message::<PrepareHeader>::try_from(
- iggy_binary_protocol::consensus::iobuf::Owned::<4096>::copy_from_slice(
- prepare.as_slice(),
- ),
- )
- .expect("partition journal entry must contain valid prepare");
+ // The partition journal already stores the wire-format
+ // `Frozen<4096>` (PrepareHeader followed by payload),
+ // so `send_to_replica` can take it directly and `clone`
+ // is a refcount bump. Matches the metadata-plane path
+ // above and avoids both the per-target 4 KiB memcpy
+ // and the prior `.expect` that would panic the shard
+ // on a corrupted journal entry.
for replica in replicas {
- send(*replica, prepare.deep_copy().into_generic()).await;
+ send(*replica, prepare.clone()).await;
}
}
}
diff --git a/core/shard/src/router.rs b/core/shard/src/router.rs
index 730d51789..c6d95aaf0 100644
--- a/core/shard/src/router.rs
+++ b/core/shard/src/router.rs
@@ -17,13 +17,14 @@
use crate::shards_table::ShardsTable;
use crate::{IggyShard, Receiver, ShardFrame};
+use crossfire::TrySendError;
use futures::FutureExt;
use iggy_binary_protocol::{
ConsensusError, ConsensusHeader, GenericHeader, Message, MessageBag, PrepareHeader,
};
use iggy_common::sharding::IggyNamespace;
use journal::{Journal, JournalHandle};
-use message_bus::MessageBus;
+use message_bus::{ConnectionInstaller, MessageBus};
use metadata::stm::StateMachine;
/// Inter-shard dispatch logic.
@@ -34,7 +35,7 @@ use metadata::stm::StateMachine;
/// pump), preventing concurrent access from independent async tasks.
impl<B, MJ, S, M, T, R> IggyShard<B, MJ, S, M, T, R>
where
- B: MessageBus,
+ B: MessageBus + ConnectionInstaller,
T: ShardsTable,
R: Send + 'static,
{
@@ -106,16 +107,29 @@ where
// headers don't carry an IggyNamespace.
0
};
- if self.senders[target as usize]
- .send(ShardFrame::fire_and_forget(generic))
- .is_err()
- {
- tracing::warn!(
- shard = self.id,
- target,
- ?operation,
- "dispatch: shard channel full or closed, message dropped"
- );
+ // `senders[target]` is a `crossfire::MTx`, which in compio is
+ // running on an io_uring reactor: a blocking `send` on a full
+ // inbox would park the reactor thread and stall every other
+ // connection on the shard. Drop instead - consensus recovers via
+ // WAL retransmit or a view change.
+ match self.senders[target as usize].try_send(ShardFrame::fire_and_forget(generic)) {
+ Ok(()) => {}
+ Err(TrySendError::Full(_)) => {
+ tracing::warn!(
+ shard = self.id,
+ target,
+ ?operation,
+ "dispatch: shard inbox full, message dropped"
+ );
+ }
+ Err(TrySendError::Disconnected(_)) => {
+ tracing::warn!(
+ shard = self.id,
+ target,
+ ?operation,
+ "dispatch: shard inbox closed, message dropped"
+ );
+ }
}
}
@@ -185,15 +199,33 @@ where
// TODO: Same view-change and Commit routing issue as dispatch() above.
0
};
- // Create a frame and send it to the target shard.
+ // Create a frame and send it to the target shard. Same non-
+ // blocking `try_send` reason as `dispatch` above: blocking on a
+ // full inbox would park the io_uring reactor.
+ //
+ // On `Full` / `Disconnected` we drop the frame and let the
+ // caller's `rx` observe the disconnect (the `response_sender`
+ // inside the frame is dropped together with the frame). Callers
+ // see an early error rather than hanging.
let (frame, rx) = ShardFrame::<R>::with_response(generic);
- if self.senders[target as usize].send(frame).is_err() {
- tracing::warn!(
- shard = self.id,
- target,
- ?operation,
- "dispatch_request: shard channel full or closed, message dropped"
- );
+ match self.senders[target as usize].try_send(frame) {
+ Ok(()) => {}
+ Err(TrySendError::Full(_)) => {
+ tracing::warn!(
+ shard = self.id,
+ target,
+ ?operation,
+ "dispatch_request: shard inbox full, message dropped"
+ );
+ }
+ Err(TrySendError::Disconnected(_)) => {
+ tracing::warn!(
+ shard = self.id,
+ target,
+ ?operation,
+ "dispatch_request: shard inbox closed, message dropped"
+ );
+ }
}
Ok(rx)
}
@@ -202,7 +234,7 @@ where
#[allow(clippy::future_not_send)]
pub async fn run_message_pump(&self, stop: Receiver<()>)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -236,7 +268,7 @@ where
#[allow(clippy::future_not_send)]
async fn process_frame(&self, frame: ShardFrame<R>)
where
- B: MessageBus<Replica = u8, Data = Message<GenericHeader>, Client = u128>,
+ B: MessageBus,
MJ: JournalHandle,
<MJ as JournalHandle>::Target: Journal<
<MJ as JournalHandle>::Storage,
@@ -249,7 +281,87 @@ where
Error = iggy_common::IggyError,
>,
{
- self.on_message(frame.message).await;
+ match frame.payload {
+ crate::ShardFramePayload::Consensus(message) => {
+ self.on_message(message).await;
+ }
+ crate::ShardFramePayload::ReplicaConnectionSetup { raw_fd, replica_id } => {
+ tracing::info!(
+ shard = self.id,
+ replica_id,
+ raw_fd,
+ "installing delegated replica fd"
+ );
+ self.bus
+ .install_replica_fd(raw_fd, replica_id, self.on_replica_message.clone());
+ }
+ crate::ShardFramePayload::ClientConnectionSetup { raw_fd, client_id } => {
+ tracing::info!(
+ shard = self.id,
+ client_id,
+ raw_fd,
+ "installing delegated client fd"
+ );
+ self.bus
+ .install_client_fd(raw_fd, client_id, self.on_client_request.clone());
+ }
+ crate::ShardFramePayload::ReplicaMappingUpdate {
+ replica_id,
+ owning_shard,
+ } => {
+ self.bus.set_shard_mapping(replica_id, owning_shard);
+ }
+ crate::ShardFramePayload::ReplicaMappingClear { replica_id } => {
+ self.bus.remove_shard_mapping(replica_id);
+ }
+ crate::ShardFramePayload::ForwardReplicaSend { replica_id, msg } => {
+ // Fast path on the owning shard: re-enter `send_to_replica`
+ // which finds the local `BusSender` in the replicas registry.
+ if let Err(e) = self.bus.send_to_replica(replica_id, msg).await {
+ tracing::debug!(
+ shard = self.id,
+ replica_id,
+ error = ?e,
+ "forward-replica-send delivery failed"
+ );
+ }
+ }
+ crate::ShardFramePayload::ForwardClientSend { client_id, msg } => {
+ if let Err(e) = self.bus.send_to_client(client_id, msg).await {
+ tracing::debug!(
+ shard = self.id,
+ client_id,
+ error = ?e,
+ "forward-client-send delivery failed"
+ );
+ }
+ }
+ crate::ShardFramePayload::ConnectionLost { replica_id } => {
+ // Shard 0 is the sole responder; the next reconnect sweep
+ // will re-dial the peer. Non-zero shards should not see this
+ // variant.
+ if self.id == 0 {
+ tracing::warn!(replica_id, "shard 0 observed replica connection loss");
+ // If a coordinator is attached, let it broadcast the
+ // clear AND drop its tracked mapping so the periodic
+ // refresh task stops re-broadcasting a dead replica's
+ // mapping. Tests that bypass the coordinator fall back
+ // to the direct broadcast.
+ if let Some(coord) = &self.coordinator {
+ coord.forget_mapping(replica_id);
+ coord.broadcast_mapping_clear(replica_id);
+ } else {
+ crate::broadcast_mapping_clear(&self.senders, replica_id);
+ }
+ } else {
+ tracing::warn!(
+ shard = self.id,
+ replica_id,
+ "non-zero shard received ConnectionLost; dropping"
+ );
+ }
+ }
+ }
// TODO: once on_message returns an R (e.g. ShardResponse), send it
// back via frame.response_sender. For now the sender is dropped and
// the caller's receiver will observe a disconnect.
diff --git a/core/simulator/src/bus.rs b/core/simulator/src/bus.rs
index 8dbd16f6e..d185e77be 100644
--- a/core/simulator/src/bus.rs
+++ b/core/simulator/src/bus.rs
@@ -15,13 +15,23 @@
// specific language governing permissions and limitations
// under the License.
+use iggy_binary_protocol::consensus::MESSAGE_ALIGN;
+use iggy_binary_protocol::consensus::iobuf::{Frozen, Owned};
use iggy_binary_protocol::{GenericHeader, Message};
-use iggy_common::IggyError;
-use message_bus::MessageBus;
+use message_bus::{MessageBus, SendError};
use std::collections::{HashSet, VecDeque};
use std::ops::Deref;
use std::sync::{Arc, Mutex};
+/// Convert an inbound Frozen (bus trait payload) into a `Message` for the
+/// simulator's packet layer. The sim keeps `Message<GenericHeader>` in its
+/// envelope because the network and tick loop reason about typed headers,
+/// not raw byte buffers. One memcpy per send is acceptable under sim.
+fn frozen_to_message(frozen: &Frozen<MESSAGE_ALIGN>) -> Message<GenericHeader> {
+ Message::try_from(Owned::<MESSAGE_ALIGN>::copy_from_slice(frozen.as_slice()))
+ .expect("simulator bus must receive a valid generic message")
+}
+
/// Message envelope for tracking sender/recipient
#[derive(Debug)]
pub enum EnvelopePayload {
@@ -69,75 +79,69 @@ impl SimOutbox {
pub fn drain(&self) -> Vec<Envelope> {
self.pending_messages.lock().unwrap().drain(..).collect()
}
-}
-impl MessageBus for SimOutbox {
- type Client = u128;
- type Replica = u8;
- type Data = Message<GenericHeader>;
- type Sender = ();
-
- fn add_client(&mut self, client: Self::Client, _sender: Self::Sender) -> bool {
- if self.clients.lock().unwrap().contains(&client) {
- return false;
- }
- self.clients.lock().unwrap().insert(client);
- true
+ /// # Panics
+ /// Panics if the internal mutex is poisoned.
+ pub fn add_client(&mut self, client: u128) -> bool {
+ self.clients.lock().unwrap().insert(client)
}
- fn remove_client(&mut self, client: Self::Client) -> bool {
- self.clients.lock().unwrap().remove(&client)
+ /// # Panics
+ /// Panics if the internal mutex is poisoned.
+ pub fn add_replica(&mut self, replica: u8) -> bool {
+ self.replicas.lock().unwrap().insert(replica)
}
- fn add_replica(&mut self, replica: Self::Replica) -> bool {
- if self.replicas.lock().unwrap().contains(&replica) {
- return false;
- }
- self.replicas.lock().unwrap().insert(replica);
- true
+ /// # Panics
+ /// Panics if the internal mutex is poisoned.
+ pub fn remove_client(&mut self, client: u128) -> bool {
+ self.clients.lock().unwrap().remove(&client)
}
- fn remove_replica(&mut self, replica: Self::Replica) -> bool {
+ /// # Panics
+ /// Panics if the internal mutex is poisoned.
+ pub fn remove_replica(&mut self, replica: u8) -> bool {
self.replicas.lock().unwrap().remove(&replica)
}
+}
+impl MessageBus for SimOutbox {
async fn send_to_client(
&self,
- client_id: Self::Client,
- message: Self::Data,
- ) -> Result<Self::Data, IggyError> {
+ client_id: u128,
+ data: Frozen<MESSAGE_ALIGN>,
+ ) -> Result<(), SendError> {
if !self.clients.lock().unwrap().contains(&client_id) {
- #[allow(clippy::cast_possible_truncation)]
- return Err(IggyError::ClientNotFound(client_id as u32));
+ return Err(SendError::ClientNotFound(client_id));
}
self.pending_messages.lock().unwrap().push_back(Envelope {
from_replica: Some(self.self_id),
to_replica: None,
to_client: Some(client_id),
- payload: EnvelopePayload::Client(message.deep_copy()),
+ payload: EnvelopePayload::Client(frozen_to_message(&data)),
});
- Ok(message)
+ Ok(())
}
async fn send_to_replica(
&self,
- replica: Self::Replica,
- message: Self::Data,
- ) -> Result<Self::Data, IggyError> {
+ replica: u8,
+ data: Frozen<MESSAGE_ALIGN>,
+ ) -> Result<(), SendError> {
if !self.replicas.lock().unwrap().contains(&replica) {
- return Err(IggyError::ResourceNotFound(format!("Replica {replica}")));
+ return Err(SendError::ReplicaNotConnected(replica));
}
self.pending_messages.lock().unwrap().push_back(Envelope {
from_replica: Some(self.self_id),
to_replica: Some(replica),
to_client: None,
- payload: EnvelopePayload::Replica(message.deep_copy()),
+ payload: EnvelopePayload::Replica(frozen_to_message(&data)),
});
- Ok(message)
+ Ok(())
}
}
@@ -153,40 +157,19 @@ impl Deref for SharedSimOutbox {
}
impl MessageBus for SharedSimOutbox {
- type Client = u128;
- type Replica = u8;
- type Data = Message<GenericHeader>;
- type Sender = ();
-
- fn add_client(&mut self, client: Self::Client, _sender: Self::Sender) -> bool {
- self.0.clients.lock().unwrap().insert(client)
- }
-
- fn remove_client(&mut self, client: Self::Client) -> bool {
- self.0.clients.lock().unwrap().remove(&client)
- }
-
- fn add_replica(&mut self, replica: Self::Replica) -> bool {
- self.0.replicas.lock().unwrap().insert(replica)
- }
-
- fn remove_replica(&mut self, replica: Self::Replica) -> bool {
- self.0.replicas.lock().unwrap().remove(&replica)
- }
-
async fn send_to_client(
&self,
- client_id: Self::Client,
- message: Self::Data,
- ) -> Result<Self::Data, IggyError> {
- self.0.send_to_client(client_id, message).await
+ client_id: u128,
+ data: Frozen<MESSAGE_ALIGN>,
+ ) -> Result<(), SendError> {
+ self.0.send_to_client(client_id, data).await
}
async fn send_to_replica(
&self,
- replica: Self::Replica,
- message: Self::Data,
- ) -> Result<Self::Data, IggyError> {
- self.0.send_to_replica(replica, message).await
+ replica: u8,
+ data: Frozen<MESSAGE_ALIGN>,
+ ) -> Result<(), SendError> {
+ self.0.send_to_replica(replica, data).await
}
}
diff --git a/core/simulator/src/lib.rs b/core/simulator/src/lib.rs
index 9cfc09b46..72b624fbb 100644
--- a/core/simulator/src/lib.rs
+++ b/core/simulator/src/lib.rs
@@ -30,7 +30,6 @@ use iggy_binary_protocol::consensus::iobuf::Owned;
use iggy_binary_protocol::{Command2, GenericHeader, Message, Operation, ReplyHeader};
use iggy_common::IggyError;
use iggy_common::sharding::IggyNamespace;
-use message_bus::MessageBus;
use network::Network;
use packet::{PacketSimulatorOptions, ProcessId};
use partitions::{Partition, PartitionOffsets, PollQueryResult, PollingArgs, PollingConsumer};
@@ -97,7 +96,7 @@ impl Simulator {
let id = i as u8;
let mut bus = SimOutbox::new(id);
for &cid in &client_ids {
- bus.add_client(cid, ());
+ bus.add_client(cid);
}
for j in 0..rc {
bus.add_replica(j);
diff --git a/core/simulator/src/replica.rs b/core/simulator/src/replica.rs
index 1a0f71133..ae1f65ea0 100644
--- a/core/simulator/src/replica.rs
+++ b/core/simulator/src/replica.rs
@@ -82,6 +82,7 @@ pub fn new_replica(id: u8, name: String, bus: &Arc<SimOutbox>, replica_count: u8
shard::IggyShard::without_inbox(
shard::ShardIdentity::new(u16::from(id), name),
+ SharedSimOutbox(Arc::clone(bus)),
metadata,
partitions,
(),