Skip to content

perf: sharded mode pipeline throughput optimization #112

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
kacy merged 3 commits into from
Feb 12, 2026
Merged

perf: sharded mode pipeline throughput optimization #112

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
648088f
perf: drain shard channel after recv to reduce select! overhead
kacy Feb 12, 2026
4814582
perf: eliminate key clones in single-key shard commands
kacy Feb 12, 2026
0ea8f8b
perf: replace join_all with dispatch-collect pipeline
kacy Feb 12, 2026
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
16 changes: 15 additions & 1 deletion crates/ember-core/src/engine.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,9 +188,23 @@ impl Engine {
}

/// Determines which shard owns a given key.
fn shard_for_key(&self, key: &str) -> usize {
pub fn shard_for_key(&self, key: &str) -> usize {
shard_index(key, self.shards.len())
}

/// Sends a request to a shard and returns the reply channel without
/// waiting for the response. Used by the connection handler to
/// dispatch commands and collect responses separately.
pub async fn dispatch_to_shard(
&self,
shard_idx: usize,
request: ShardRequest,
) -> Result<tokio::sync::oneshot::Receiver<ShardResponse>, ShardError> {
if shard_idx >= self.shards.len() {
return Err(ShardError::Unavailable);
}
self.shards[shard_idx].dispatch(request).await
}
}

/// Pure function: maps a key to a shard index.
Expand Down
150 changes: 97 additions & 53 deletions crates/ember-core/src/shard.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -464,8 +464,9 @@ impl ShardHandle {

/// Sends a request and returns the reply channel without waiting
/// for the response. Used by `Engine::broadcast` to fan out to
/// all shards before collecting results.
pub(crate) async fn dispatch(
/// all shards before collecting results, and by
/// `Engine::dispatch_to_shard` for the dispatch-collect pipeline.
pub async fn dispatch(
&self,
request: ShardRequest,
) -> Result<oneshot::Receiver<ShardResponse>, ShardError> {
Expand Down Expand Up @@ -651,63 +652,35 @@ async fn run_shard(
msg = rx.recv() => {
match msg {
Some(msg) => {
let request_kind = describe_request(&msg.request);
let response = dispatch(
process_message(
msg,
&mut keyspace,
&msg.request,
&mut aof_writer,
fsync_policy,
&persistence,
&drop_handle,
shard_id,
#[cfg(feature = "protobuf")]
&schema_registry,
);

// write AOF record for successful mutations
if let Some(ref mut writer) = aof_writer {
if let Some(record) = to_aof_record(&msg.request, &response) {
if let Err(e) = writer.write_record(&record) {
warn!(shard_id, "aof write failed: {e}");
}
if fsync_policy == FsyncPolicy::Always {
if let Err(e) = writer.sync() {
warn!(shard_id, "aof sync failed: {e}");
}
}
}
}

// handle snapshot/rewrite (these need mutable access
// to both keyspace and aof_writer)
match request_kind {
RequestKind::Snapshot => {
let resp = handle_snapshot(
&keyspace, &persistence, shard_id,
);
let _ = msg.reply.send(resp);
continue;
}
RequestKind::RewriteAof => {
let resp = handle_rewrite(
&keyspace,
&persistence,
&mut aof_writer,
shard_id,
#[cfg(feature = "protobuf")]
&schema_registry,
);
let _ = msg.reply.send(resp);
continue;
}
RequestKind::FlushDbAsync => {
let old_entries = keyspace.flush_async();
if let Some(ref handle) = drop_handle {
handle.defer_entries(old_entries);
}
// else: old_entries drops inline here
let _ = msg.reply.send(ShardResponse::Ok);
continue;
}
RequestKind::Other => {}
// drain any pending messages without re-entering select!.
// this amortizes the select! overhead across bursts of
// pipelined commands that arrived while we processed the
// first message.
while let Ok(msg) = rx.try_recv() {
process_message(
msg,
&mut keyspace,
&mut aof_writer,
fsync_policy,
&persistence,
&drop_handle,
shard_id,
#[cfg(feature = "protobuf")]
&schema_registry,
);
}

let _ = msg.reply.send(response);
}
None => break, // channel closed, shard shutting down
}
Expand All @@ -731,6 +704,77 @@ async fn run_shard(
}
}

/// Processes a single shard message: dispatches the command, writes
/// the AOF record, handles special requests, and sends the reply.
///
/// Extracted from the main loop so it can be called for both the
/// initial `recv()` and the `try_recv()` drain loop.
#[allow(clippy::too_many_arguments)]
fn process_message(
msg: ShardMessage,
keyspace: &mut Keyspace,
aof_writer: &mut Option<AofWriter>,
fsync_policy: FsyncPolicy,
persistence: &Option<ShardPersistenceConfig>,
drop_handle: &Option<DropHandle>,
shard_id: u16,
#[cfg(feature = "protobuf")] schema_registry: &Option<crate::schema::SharedSchemaRegistry>,
) {
let request_kind = describe_request(&msg.request);
let response = dispatch(
keyspace,
&msg.request,
#[cfg(feature = "protobuf")]
schema_registry,
);

// write AOF record for successful mutations
if let Some(ref mut writer) = aof_writer {
if let Some(record) = to_aof_record(&msg.request, &response) {
if let Err(e) = writer.write_record(&record) {
warn!(shard_id, "aof write failed: {e}");
}
if fsync_policy == FsyncPolicy::Always {
if let Err(e) = writer.sync() {
warn!(shard_id, "aof sync failed: {e}");
}
}
}
}

// handle special requests that need access to persistence state
match request_kind {
RequestKind::Snapshot => {
let resp = handle_snapshot(keyspace, persistence, shard_id);
let _ = msg.reply.send(resp);
return;
}
RequestKind::RewriteAof => {
let resp = handle_rewrite(
keyspace,
persistence,
aof_writer,
shard_id,
#[cfg(feature = "protobuf")]
schema_registry,
);
let _ = msg.reply.send(resp);
return;
}
RequestKind::FlushDbAsync => {
let old_entries = keyspace.flush_async();
if let Some(ref handle) = drop_handle {
handle.defer_entries(old_entries);
}
let _ = msg.reply.send(ShardResponse::Ok);
return;
}
RequestKind::Other => {}
}

let _ = msg.reply.send(response);
}

/// Lightweight tag so we can identify requests that need special
/// handling after dispatch without borrowing the request again.
enum RequestKind {
Expand Down
Loading
Loading