Skip to content

fix(shard): bounded cross-shard fan-out + reply awaits (c10k E1/E3/E4), lazy rehydrated buffers (D3) #441

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
TinDang97 merged 1 commit into from
Aug 7, 2026
Merged

fix(shard): bounded cross-shard fan-out + reply awaits (c10k E1/E3/E4), lazy rehydrated buffers (D3) #441

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

## [Unreleased]

### Fixed
- **Cross-shard fan-out no longer drops messages silently on a full SPSC ring
(c10k E1/E3), and cross-shard reply awaits are bounded (E4).** PUBLISH
fan-out (immediate, batched, and EXEC-queued) and SCRIPT LOAD propagation
used a single `try_push` — a transiently-full ring lost the message with no
log or metric: subscribers on that shard silently missed the publish, or its
script cache diverged (NOSCRIPT for a sha the server had just returned).
All five sites now retry with the same bounded, shutdown-aware backpressure
as the command dispatch path, and a final give-up is loud
(`moon_xshard_fanout_drop_total`). Reply awaits on the slotted dispatch and
publish paths — previously unbounded, so one wedged shard could hang a
client task forever — now share the coordinator's 30 s bound: publish counts
degrade (under-report, counted by `moon_xshard_reply_timeout_total`), while
a slotted-dispatch timeout errors the batch and closes the connection,
because the per-connection reply slot cannot be safely reused after an
abandoned await. Rehydrated connection handlers (park wake / migration)
also start with 512 B I/O buffers instead of 3×8 KiB, shrinking the memory
spike of fleet-synchronized wakes; buffers grow back on first real traffic
(c10k D3).

### Added
- **`INFO memory` can now explain a `used_memory`-vs-RSS gap.** Build with
`--features jemalloc-stats` and `INFO memory` reports Redis's `allocator_*`
Expand Down
25 changes: 25 additions & 0 deletions src/admin/metrics_setup.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -923,6 +923,31 @@ pub fn record_dispatch_cross_spsc() {
counter!("moon_dispatch_path_total", "path" => "cross_spsc").increment(1);
}

/// Cross-shard fan-out message dropped after bounded retry (c10k E1/E3):
/// the target shard's SPSC ring stayed full through every backoff. `kind` is
/// `"publish"` (that shard's subscribers miss the message; PUBLISH count
/// under-reports) or `"script_load"` (that shard's script cache diverges —
/// EVALSHA there answers NOSCRIPT until the next SCRIPT LOAD).
#[inline]
pub fn record_xshard_fanout_drop(kind: &'static str) {
if !METRICS_INITIALIZED.load(Ordering::Relaxed) {
return;
}
counter!("moon_xshard_fanout_drop_total", "kind" => kind).increment(1);
}

/// Cross-shard reply await expired (c10k E4): the owner shard did not fill
/// the reply slot within `XSHARD_REPLY_TIMEOUT`. `kind` `"dispatch"` is
/// fatal for the connection (the reusable slot may be filled late);
/// `"publish"` degrades to an under-reported subscriber count.
#[inline]
pub fn record_xshard_reply_timeout(kind: &'static str) {
if !METRICS_INITIALIZED.load(Ordering::Relaxed) {
return;
}
counter!("moon_xshard_reply_timeout_total", "kind" => kind).increment(1);
}

/// Batched variant of `record_dispatch_cross_spsc`.
#[inline]
pub fn record_dispatch_cross_spsc_batch(count: u64) {
Expand Down
28 changes: 7 additions & 21 deletions src/server/conn/handler_monoio/dispatch.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
//! Each helper returns `true` if the command was consumed (caller should `continue`).

use bytes::Bytes;
use ringbuf::traits::Producer;
use std::cell::RefCell;
use std::rc::Rc;
use std::sync::Arc;
Expand All @@ -20,8 +19,6 @@ use crate::runtime::cancel::CancellationToken;
use crate::runtime::channel;
use crate::server::conn::core::{ConnectionContext, ConnectionState};
use crate::server::conn::util::extract_bytes;
use crate::shard::dispatch::ShardMessage;
use crate::shard::mesh::ChannelMesh;
use crate::tracking::TrackingState;
use crate::workspace::strip_workspace_prefix_from_response;

Expand Down Expand Up @@ -218,12 +215,14 @@ pub(super) fn try_handle_eval(

/// Handle SCRIPT subcommands (LOAD, EXISTS, FLUSH). Returns `true` if consumed.
///
/// `#[inline]`: see `try_handle_cluster` rationale.
#[inline]
pub(super) fn try_handle_script(
/// Async since E3: the SCRIPT LOAD fan-out retries a full ring with bounded
/// backpressure instead of silently dropping the load (divergent per-shard
/// script caches). Cold path — SCRIPT is never hot.
pub(super) async fn try_handle_script(
cmd: &[u8],
cmd_args: &[Frame],
ctx: &ConnectionContext,
shutdown: &crate::runtime::cancel::CancellationToken,
responses: &mut Vec<Frame>,
) -> bool {
if !cmd.eq_ignore_ascii_case(b"SCRIPT") {
Expand All @@ -232,21 +231,8 @@ pub(super) fn try_handle_script(
let (response, fanout) =
crate::scripting::handle_script_subcommand(&ctx.script_cache, cmd_args);
if let Some((sha1, script_bytes)) = fanout {
let mut producers = ctx.dispatch_tx.borrow_mut();
for target in 0..ctx.num_shards {
if target == ctx.shard_id {
continue;
}
let idx = ChannelMesh::target_index(ctx.shard_id, target);
let msg = ShardMessage::ScriptLoad {
sha1: sha1.clone(),
script: script_bytes.clone(),
};
if producers[idx].try_push(msg).is_ok() {
ctx.spsc_notifiers[target].notify_one();
}
}
drop(producers);
crate::server::conn::shared::script_fanout_bounded(ctx, shutdown, &sha1, &script_bytes)
.await;
}
responses.push(response);
true
Expand Down
173 changes: 139 additions & 34 deletions src/server/conn/handler_monoio/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -333,14 +333,23 @@ pub(crate) async fn handle_connection_sharded_monoio<
// NOTE: do NOT call record_connection_opened() here — the caller
// (conn_accept.rs) already increments via try_accept_connection().

// c10k D3 (lazy resumed-buffer sizing): a fresh connection warms up with
// the full 8 KiB per buffer (no growth stall on its first pipeline). A
// REHYDRATED handler (park wake / migration) starts at 512 B: wakes
// arrive in fleet-sized bursts (synchronized keepalives), and 2×8 KiB
// per wake dominated the burst working set before the first byte was
// even parsed. BytesMut grows on demand — a busy resumed conn pays one
// amortized regrow; the idle majority never pays the 16 KiB at all.
let rehydrated = migrated_state.is_some();
let init_cap = if rehydrated { 512 } else { 8192 };
let mut read_buf = if initial_read_buf.is_empty() {
BytesMut::with_capacity(8192)
BytesMut::with_capacity(init_cap)
} else {
let mut buf = initial_read_buf;
buf.reserve(8192);
buf.reserve(init_cap);
buf
};
let mut write_buf = BytesMut::with_capacity(8192);
let mut write_buf = BytesMut::with_capacity(init_cap);
// c10k A1: set when a blocking command leaves unparsed input in
// `read_buf` (see the read-skip guard in the main loop).
let mut carried_input = false;
Expand Down Expand Up @@ -401,7 +410,10 @@ pub(crate) async fn handle_connection_sharded_monoio<

// Pre-allocate read buffer outside the loop to avoid per-read heap allocation.
// Monoio's ownership I/O takes ownership and returns the buffer, so we reassign.
let mut tmp_buf = vec![0u8; 8192];
// D3: same lazy sizing as read_buf/write_buf — the first read of a
// rehydrated conn is usually a probe-sized frame (keepalive PING); the
// shrink logic at the loop tail governs the steady state either way.
let mut tmp_buf = vec![0u8; init_cap];

// c10k W11: two-stage idle park (see idle_park.rs). Cancel-capable
// streams register for the shard chore's ≥1s sweep; `downshifted` tracks
Expand Down Expand Up @@ -980,6 +992,15 @@ pub(crate) async fn handle_connection_sharded_monoio<
}
}

// D3: a rehydrated conn starts with a small (512 B) owned read buffer;
// the moment a read saturates it (real traffic, not a keepalive-sized
// probe), restore the full 8 KiB so bulk transfers aren't capped at
// 512 B per syscall. Sited with the C2 check below: after every read
// arm, once per iteration.
if tmp_buf.len() < 8192 && read_buf.len() >= tmp_buf.len() {
tmp_buf = vec![0u8; 8192];
}

// c10k C2: query-buffer ceiling. One check per read iteration, sited
// after every read arm and ahead of both parse paths — an incomplete
// frame is exactly what makes `read_buf` grow, and an incomplete
Expand Down Expand Up @@ -1282,7 +1303,9 @@ pub(crate) async fn handle_connection_sharded_monoio<
{
continue;
}
if cmd_len == 6 && dispatch::try_handle_script(cmd, cmd_args, ctx, &mut responses) {
if cmd_len == 6
&& dispatch::try_handle_script(cmd, cmd_args, ctx, &shutdown, &mut responses).await
{
Comment on lines +1306 to +1308

@qodo-code-review qodo-code-review Bot Aug 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Informational

2. Async script check on length 🐞 Bug ➹ Performance

The monoio handler now awaits try_handle_script for every 6-byte command, even when the command is
not SCRIPT (the async function immediately returns false). This adds avoidable async-future
construction/poll overhead to common 6-letter commands (e.g., CONFIG/SELECT/GETSET).
Agent Prompt 
### Issue description
`handler_monoio/mod.rs` calls/awaits an async `try_handle_script` based only on `cmd_len == 6`. For all other 6-byte commands, `try_handle_script` immediately returns `false` after its internal `eq_ignore_ascii_case(b"SCRIPT")` check, so the await is unnecessary work on a hot path.

### Issue Context
`try_handle_script` became `async` to support bounded fan-out retries for `SCRIPT LOAD` (cold path), but the current call-site pattern makes the async wrapper run for unrelated 6-byte commands.

### Fix Focus Areas
- src/server/conn/handler_monoio/mod.rs[1294-1330]
- src/server/conn/handler_monoio/dispatch.rs[216-239]

### Suggested fix
- Change the call site to avoid invoking the async helper unless the command is actually SCRIPT, e.g.:
  - `if cmd_len == 6 && cmd.eq_ignore_ascii_case(b"SCRIPT") { if dispatch::try_handle_script(...).await { continue; } }`
- Or split into:
  - a sync predicate (`is_script(cmd) -> bool`), and
  - an async handler only executed on true.

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

continue;
}
if dispatch::try_handle_cluster_routing(cmd, cmd_args, &mut conn, ctx, &mut responses) {
Expand Down Expand Up @@ -1548,7 +1571,10 @@ pub(crate) async fn handle_connection_sharded_monoio<
) {
Some(err) => err,
None => Frame::Integer(
crate::server::conn::shared::publish_post_txn(ctx, &ch, &msg).await,
crate::server::conn::shared::publish_post_txn(
ctx, &shutdown, &ch, &msg,
)
.await,
),
};
if let Frame::Array(items) = &mut responses[exec_idx] {
Expand Down Expand Up @@ -2544,33 +2570,75 @@ pub(crate) async fn handle_connection_sharded_monoio<
std::sync::Arc<crate::shard::dispatch::PubSubResponseSlot>,
Vec<usize>,
)> = Vec::new();
{
let mut producers = ctx.dispatch_tx.borrow_mut();
for (target, entries) in publish_batches.drain() {
let n = entries.len();
let slot = std::sync::Arc::new(
crate::shard::dispatch::PubSubResponseSlot::with_counts(1, n),
);
let resp_indices: Vec<usize> = entries.iter().map(|(idx, _, _)| *idx).collect();
let pairs: Vec<(Bytes, Bytes)> =
entries.into_iter().map(|(_, ch, msg)| (ch, msg)).collect();

let idx = ChannelMesh::target_index(ctx.shard_id, target);
let batch_msg = ShardMessage::PubSubPublishBatch {
pairs,
slot: slot.clone(),
};
if producers[idx].try_push(batch_msg).is_ok() {
for (target, entries) in publish_batches.drain() {
let n = entries.len();
let slot = std::sync::Arc::new(
crate::shard::dispatch::PubSubResponseSlot::with_counts(1, n),
);
let resp_indices: Vec<usize> = entries.iter().map(|(idx, _, _)| *idx).collect();
let pairs: Vec<(Bytes, Bytes)> =
entries.into_iter().map(|(_, ch, msg)| (ch, msg)).collect();

let idx = ChannelMesh::target_index(ctx.shard_id, target);
// E1: bounded backpressure retry instead of one bare try_push
// — a transiently-full ring no longer loses the batch. Borrow
// taken+released per attempt, never held across the backoff
// await (tokio parity — handler_sharded).
let mut pending = Some(ShardMessage::PubSubPublishBatch {
pairs,
slot: slot.clone(),
});
let outcome = crate::shard::dispatch::push_with_backpressure(
&shutdown,
crate::shard::dispatch::CROSS_SHARD_PUSH_MAX_RETRIES,
crate::shard::dispatch::CROSS_SHARD_PUSH_BACKOFF,
|| match pending.take() {
None => true,
Some(m) => {
let mut producers = ctx.dispatch_tx.borrow_mut();
match producers[idx].try_push(m) {
Ok(()) => true,
Err(back) => {
pending = Some(back);
false
}
}
}
},
)
.await;
match outcome {
crate::shard::dispatch::PushOutcome::Pushed => {
ctx.spsc_notifiers[target].notify_one();
} else {
slot.add(0); // push failed, mark as done
}
batch_slots.push((slot, resp_indices));
outcome => {
// Give-up: deliver-to-zero so the reply can't hang —
// but loudly (was a silent drop pre-E1).
tracing::warn!(
"Shard {}: PUBLISH batch fan-out to shard {target} dropped ({outcome:?})",
ctx.shard_id
);
crate::admin::metrics_setup::record_xshard_fanout_drop("publish");
slot.add(0);
}
}
batch_slots.push((slot, resp_indices));
}
// Resolve all batch slots
// Resolve all batch slots (E4: bounded — a wedged shard degrades
// to an under-reported count, never a hung client).
for (slot, resp_indices) in &batch_slots {
crate::shard::dispatch::PubSubResponseFuture::new(slot.clone()).await;
if !crate::shard::dispatch::await_pubsub_slot_bounded(
slot,
crate::shard::dispatch::XSHARD_REPLY_TIMEOUT,
)
.await
{
tracing::warn!(
"Shard {}: PUBLISH batch reply timed out awaiting remote shard",
ctx.shard_id
);
crate::admin::metrics_setup::record_xshard_reply_timeout("publish");
}
for (i, resp_idx) in resp_indices.iter().enumerate() {
let remote_count = slot.counts[i].load(std::sync::atomic::Ordering::Relaxed);
if remote_count > 0 {
Expand All @@ -2582,6 +2650,12 @@ pub(crate) async fn handle_connection_sharded_monoio<
}
}

// E4: set when a cross-shard reply await times out. The per-connection
// ResponseSlot is REUSED across batches — a late fill after a timeout
// would be read by the NEXT batch as its own reply — so a timeout is
// fatal: error the affected entries, flush, then close the connection.
let mut xshard_reply_fatal = false;

// Phase 2b: Dispatch all deferred remote commands as batched
// PipelineBatchSlotted messages (one per target shard), await all in parallel.
if !remote_groups.is_empty() {
Expand Down Expand Up @@ -2674,10 +2748,11 @@ pub(crate) async fn handle_connection_sharded_monoio<
// keeps the slot alive until the last handle drops. (This replaced the
// old raw-pointer-into-stack-pool design, whose contract required the
// await to run to completion to avoid a panic-unwind UAF; see the
// `ResponseSlotPtr` doc + PR review.) The await is still unbounded here
// for simplicity — the target shard always fills every message it
// drains — but a shutdown-aware bound is now a safe, tracked follow-up.
// The tokio handler carries the identical await.
// `ResponseSlotPtr` doc + PR review.) Since E4 the await is BOUNDED
// (XSHARD_REPLY_TIMEOUT) — expiry errors the batch and closes the
// connection, because the reused slot could otherwise hand a late
// fill to the next batch. The tokio handler carries the identical
// bounded await.
//
// C2 pipeline guard (see XSHARD_SPIN_MAX_BATCH_REMOTE): total cross-shard
// commands in THIS batch. The reply-side spin may engage only for a singleton
Expand Down Expand Up @@ -2705,9 +2780,33 @@ pub(crate) async fn handle_connection_sharded_monoio<
}
}
match spun {
Some(r) => r,
None => response_pool.future_for(target).await,
Some(r) => Some(r),
// E4: bounded await — a wedged owner shard can no
// longer hang this client task forever.
None => {
crate::shard::dispatch::await_response_slot_bounded(
response_pool.future_for(target),
crate::shard::dispatch::XSHARD_REPLY_TIMEOUT,
)
.await
}
}
};
let Some(shard_responses) = shard_responses else {
tracing::error!(
"Shard {}: cross-shard reply from shard {target} timed out; \
failing the batch and closing the connection (slot unsafe to reuse)",
ctx.shard_id
);
crate::admin::metrics_setup::record_xshard_reply_timeout("dispatch");
for (resp_idx, _, _, _) in &meta {
responses[*resp_idx] =
Frame::Error(Bytes::from_static(b"ERR cross-shard reply timeout"));
}
xshard_reply_fatal = true;
// Skip the fsync barrier too: with no reply there is
// nothing durable to confirm for these entries.
continue;
};
// H1-BARRIER: collect write resp_idxs before consuming meta
// so we can overwrite them if the fsync barrier fails.
Expand Down Expand Up @@ -2800,6 +2899,12 @@ pub(crate) async fn handle_connection_sharded_monoio<
}
}

// E4: a timed-out cross-shard reply slot must never be reused — the
// error replies are flushed above, now close.
if xshard_reply_fatal {
break;
}

// Update live state after each batch — lock-free (QW8, 2026-06
// review: this was a global registry write lock per batch), and
// clock-free (shard-cached ms, not Instant::now()).
Expand Down
Loading
Loading