Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
41 changes: 41 additions & 0 deletions crates/ember-protocol/src/command.rs
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,15 @@ pub enum Command {
/// ASKING. Signals that the next command is for a migrating slot.
Asking,

/// SLOWLOG GET [count]. Returns recent slow log entries.
SlowLogGet { count: Option<usize> },

/// SLOWLOG LEN. Returns the number of entries in the slow log.
SlowLogLen,

/// SLOWLOG RESET. Clears the slow log.
SlowLogReset,

/// A command we don't recognize (yet).
Unknown(String),
}
Expand Down Expand Up @@ -359,6 +368,9 @@ impl Command {
Command::ClusterGetKeysInSlot { .. } => "cluster_getkeysinslot",
Command::Migrate { .. } => "migrate",
Command::Asking => "asking",
Command::SlowLogGet { .. } => "slowlog",
Command::SlowLogLen => "slowlog",
Command::SlowLogReset => "slowlog",
Command::Unknown(_) => "unknown",
}
}
Expand Down Expand Up @@ -439,6 +451,7 @@ impl Command {
"CLUSTER" => parse_cluster(&frames[1..]),
"ASKING" => parse_asking(&frames[1..]),
"MIGRATE" => parse_migrate(&frames[1..]),
"SLOWLOG" => parse_slowlog(&frames[1..]),
_ => Ok(Command::Unknown(name)),
}
}
Expand Down Expand Up @@ -1298,6 +1311,34 @@ fn parse_asking(args: &[Frame]) -> Result<Command, ProtocolError> {
Ok(Command::Asking)
}

fn parse_slowlog(args: &[Frame]) -> Result<Command, ProtocolError> {
if args.is_empty() {
return Err(ProtocolError::WrongArity("SLOWLOG".into()));
}

let subcmd = extract_string(&args[0])?.to_ascii_uppercase();
match subcmd.as_str() {
"GET" => {
let count = if args.len() > 1 {
let n: usize = extract_string(&args[1])?
.parse()
.map_err(|_| {
ProtocolError::InvalidCommandFrame("invalid count for SLOWLOG GET".into())
})?;
Some(n)
} else {
None
};
Ok(Command::SlowLogGet { count })
}
"LEN" => Ok(Command::SlowLogLen),
"RESET" => Ok(Command::SlowLogReset),
other => Err(ProtocolError::InvalidCommandFrame(format!(
"unknown SLOWLOG subcommand '{other}'"
))),
}
}

fn parse_slot_list(args: &[Frame]) -> Result<Vec<u16>, ProtocolError> {
let mut slots = Vec::with_capacity(args.len());
for arg in args {
Expand Down
243 changes: 202 additions & 41 deletions crates/ember-server/src/connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
//! sharded engine, and writes responses back. Supports pipelining
//! by processing multiple frames from a single read.

use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::{Duration, Instant};

use bytes::{Bytes, BytesMut};
Expand All @@ -12,6 +14,9 @@ use ember_protocol::{parse_frame, Command, Frame, SetExpire};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

use crate::server::ServerContext;
use crate::slowlog::SlowLog;

/// Initial read buffer capacity. 4KB covers most commands comfortably
/// without over-allocating for simple PING/SET/GET workloads.
const BUF_CAPACITY: usize = 4096;
Expand All @@ -33,6 +38,8 @@ const IDLE_TIMEOUT: Duration = Duration::from_secs(300); // 5 minutes
pub async fn handle(
mut stream: TcpStream,
engine: Engine,
ctx: &Arc<ServerContext>,
slow_log: &Arc<SlowLog>,
metrics_enabled: bool,
) -> Result<(), Box<dyn std::error::Error>> {
// disable Nagle's algorithm — cache servers need low-latency writes,
Expand Down Expand Up @@ -68,7 +75,8 @@ pub async fn handle(
match parse_frame(&buf) {
Ok(Some((frame, consumed))) => {
let _ = buf.split_to(consumed);
let response = process(frame, &engine, metrics_enabled).await;
let response =
process(frame, &engine, ctx, slow_log, metrics_enabled).await;
response.serialize(&mut out);
}
Ok(None) => break, // need more data
Expand All @@ -91,21 +99,27 @@ pub async fn handle(
///
/// When `metrics_enabled` is true, records per-command latency and
/// error counters in prometheus.
async fn process(frame: Frame, engine: &Engine, metrics_enabled: bool) -> Frame {
async fn process(
frame: Frame,
engine: &Engine,
ctx: &Arc<ServerContext>,
slow_log: &Arc<SlowLog>,
metrics_enabled: bool,
) -> Frame {
match Command::from_frame(frame) {
Ok(cmd) => {
let cmd_name = cmd.command_name();
let start = if metrics_enabled {
Some(Instant::now())
} else {
None
};
let start = Instant::now();

let response = execute(cmd, engine).await;
let response = execute(cmd, engine, ctx, slow_log).await;
let elapsed = start.elapsed();

if let Some(start) = start {
ctx.commands_processed.fetch_add(1, Ordering::Relaxed);
slow_log.maybe_record(elapsed, cmd_name);

if metrics_enabled {
let is_error = matches!(&response, Frame::Error(_));
crate::metrics::record_command(cmd_name, start.elapsed(), is_error);
crate::metrics::record_command(cmd_name, elapsed, is_error);
}

response
Expand All @@ -119,7 +133,12 @@ async fn process(frame: Frame, engine: &Engine, metrics_enabled: bool) -> Frame
/// Ping and Echo are handled inline (no shard routing needed).
/// Single-key commands route to the owning shard. Multi-key commands
/// (DEL, EXISTS) fan out across shards and aggregate results.
async fn execute(cmd: Command, engine: &Engine) -> Frame {
async fn execute(
cmd: Command,
engine: &Engine,
ctx: &Arc<ServerContext>,
slow_log: &Arc<SlowLog>,
) -> Frame {
match cmd {
// -- no shard needed --
Command::Ping(None) => Frame::Simple("PONG".into()),
Expand Down Expand Up @@ -330,36 +349,7 @@ async fn execute(cmd: Command, engine: &Engine) -> Frame {
},

Command::Info { section } => {
let section_upper = section.as_deref().map(|s| s.to_ascii_uppercase());
match section_upper.as_deref() {
None | Some("KEYSPACE") => match engine.broadcast(|| ShardRequest::Stats).await {
Ok(responses) => {
let mut total = KeyspaceStats {
key_count: 0,
used_bytes: 0,
keys_with_expiry: 0,
keys_expired: 0,
keys_evicted: 0,
};
for r in &responses {
if let ShardResponse::Stats(stats) = r {
total.key_count += stats.key_count;
total.used_bytes += stats.used_bytes;
total.keys_with_expiry += stats.keys_with_expiry;
total.keys_expired += stats.keys_expired;
total.keys_evicted += stats.keys_evicted;
}
}
let info = format!(
"# Keyspace\r\ndb0:keys={},expires={},used_bytes={}\r\n",
total.key_count, total.keys_with_expiry, total.used_bytes
);
Frame::Bulk(Bytes::from(info))
}
Err(e) => Frame::Error(format!("ERR {e}")),
},
Some(other) => Frame::Error(format!("ERR unsupported INFO section '{other}'")),
}
render_info(engine, ctx, section.as_deref()).await
}

Command::BgSave => match engine.broadcast(|| ShardRequest::Snapshot).await {
Expand Down Expand Up @@ -915,6 +905,35 @@ async fn execute(cmd: Command, engine: &Engine) -> Frame {
Frame::Error("ERR This instance has cluster support disabled".into())
}

// -- slow log commands --
Command::SlowLogGet { count } => {
let entries = slow_log.get(count);
let frames: Vec<Frame> = entries
.into_iter()
.map(|e| {
let ts = e
.timestamp
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
Frame::Array(vec![
Frame::Integer(e.id as i64),
Frame::Integer(ts as i64),
Frame::Integer(e.duration.as_micros() as i64),
Frame::Array(vec![Frame::Bulk(Bytes::from(e.command))]),
])
})
.collect();
Frame::Array(frames)
}

Command::SlowLogLen => Frame::Integer(slow_log.len() as i64),

Command::SlowLogReset => {
slow_log.reset();
Frame::Simple("OK".into())
}

Command::Unknown(name) => Frame::Error(format!("ERR unknown command '{name}'")),
}
}
Expand All @@ -940,6 +959,148 @@ where
}
}

/// Renders the INFO response with multiple sections.
///
/// With no argument, returns all sections. With a section name,
/// returns only that section. Matches Redis convention of `#` headers
/// followed by `key:value` pairs separated by `\r\n`.
async fn render_info(
engine: &Engine,
ctx: &Arc<ServerContext>,
section: Option<&str>,
) -> Frame {
let section_upper = section.map(|s| s.to_ascii_uppercase());
let want_all = section_upper.is_none();
let want = |name: &str| want_all || section_upper.as_deref() == Some(name);

// only broadcast to shards if we need keyspace/memory/persistence sections
let stats = if want("KEYSPACE") || want("MEMORY") || want("PERSISTENCE") || want("STATS") {
match engine.broadcast(|| ShardRequest::Stats).await {
Ok(responses) => {
let mut total = KeyspaceStats {
key_count: 0,
used_bytes: 0,
keys_with_expiry: 0,
keys_expired: 0,
keys_evicted: 0,
};
for r in &responses {
if let ShardResponse::Stats(s) = r {
total.key_count += s.key_count;
total.used_bytes += s.used_bytes;
total.keys_with_expiry += s.keys_with_expiry;
total.keys_expired += s.keys_expired;
total.keys_evicted += s.keys_evicted;
}
}
Some(total)
}
Err(e) => return Frame::Error(format!("ERR {e}")),
}
} else {
None
};

let mut out = String::with_capacity(512);

if want("SERVER") {
let uptime = ctx.start_time.elapsed().as_secs();
out.push_str("# Server\r\n");
out.push_str(&format!("ember_version:{}\r\n", ctx.version));
out.push_str(&format!("process_id:{}\r\n", std::process::id()));
out.push_str(&format!("uptime_in_seconds:{uptime}\r\n"));
out.push_str(&format!("shard_count:{}\r\n", ctx.shard_count));
out.push_str("\r\n");
}

if want("CLIENTS") {
let connected = ctx.connections_accepted.load(Ordering::Relaxed);
out.push_str("# Clients\r\n");
out.push_str(&format!("connected_clients:{connected}\r\n"));
out.push_str(&format!("max_clients:{}\r\n", ctx.max_connections));
out.push_str("\r\n");
}

if want("MEMORY") {
if let Some(ref stats) = stats {
out.push_str("# Memory\r\n");
out.push_str(&format!("used_memory:{}\r\n", stats.used_bytes));
out.push_str(&format!(
"used_memory_human:{}\r\n",
human_bytes(stats.used_bytes)
));
if let Some(max) = ctx.max_memory {
out.push_str(&format!("max_memory:{max}\r\n"));
out.push_str(&format!("max_memory_human:{}\r\n", human_bytes(max)));
} else {
out.push_str("max_memory:0\r\n");
out.push_str("max_memory_human:unlimited\r\n");
}
out.push_str("\r\n");
}
}

if want("PERSISTENCE") {
out.push_str("# Persistence\r\n");
out.push_str(&format!(
"aof_enabled:{}\r\n",
if ctx.aof_enabled { 1 } else { 0 }
));
out.push_str("\r\n");
}

if want("STATS") {
let total_conns = ctx.connections_accepted.load(Ordering::Relaxed);
let total_cmds = ctx.commands_processed.load(Ordering::Relaxed);
out.push_str("# Stats\r\n");
out.push_str(&format!("total_connections_received:{total_conns}\r\n"));
out.push_str(&format!("total_commands_processed:{total_cmds}\r\n"));
if let Some(ref stats) = stats {
out.push_str(&format!("expired_keys:{}\r\n", stats.keys_expired));
out.push_str(&format!("evicted_keys:{}\r\n", stats.keys_evicted));
}
out.push_str("\r\n");
}

if want("KEYSPACE") {
if let Some(ref stats) = stats {
out.push_str("# Keyspace\r\n");
if stats.key_count > 0 {
out.push_str(&format!(
"db0:keys={},expires={},used_bytes={}\r\n",
stats.key_count, stats.keys_with_expiry, stats.used_bytes
));
}
out.push_str("\r\n");
}
}

// trim trailing blank line
if out.ends_with("\r\n\r\n") {
out.truncate(out.len() - 2);
}

Frame::Bulk(Bytes::from(out))
}

/// Formats a byte count as a human-readable string (e.g. "1.23M").
fn human_bytes(bytes: usize) -> String {
const KB: f64 = 1024.0;
const MB: f64 = KB * 1024.0;
const GB: f64 = MB * 1024.0;

let b = bytes as f64;
if b >= GB {
format!("{:.2}G", b / GB)
} else if b >= MB {
format!("{:.2}M", b / MB)
} else if b >= KB {
format!("{:.2}K", b / KB)
} else {
format!("{bytes}B")
}
}

/// Returns the standard WRONGTYPE error frame.
fn wrongtype_error() -> Frame {
Frame::Error("WRONGTYPE Operation against a key holding the wrong kind of value".into())
Expand Down
Loading
Loading