From ce31bf68f0d97182367e01a359be7a70c7a427d5 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Wed, 25 Feb 2026 12:39:01 -0500 Subject: [PATCH 1/5] feat(client): typed API and pipelining for ember-client MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit adds a full typed command API and pipeline builder on top of the existing raw `send()` method. no new crate — the types layer cleanly on what's there. commands covered: get/set/set_ex/del/exists/expire/persist/ttl/pttl/ incr/decr/incrby/decrby/append/mget/mset/getdel (strings), lpush/rpush/ lpop/rpop/lrange/llen (lists), hset/hget/hgetall/hdel/hexists/hlen/hkeys/ hvals (hashes), sadd/srem/smembers/sismember/scard (sets), zadd/zrange/ zrange_withscores/zscore/zrank/zrem/zcard (sorted sets), ping/dbsize/flushdb. `Pipeline` is a chainable builder that writes all frames in a single syscall and reads responses sequentially, keeping pipelining cost to one flush per batch. `ClientError::Server` is new — server-returned errors (`WRONGTYPE`, etc.) are now distinct from wire-level `Protocol` errors. --- Cargo.lock | 1 + crates/ember-client/README.md | 96 +++ crates/ember-client/src/commands.rs | 896 ++++++++++++++++++++++ crates/ember-client/src/connection.rs | 44 +- crates/ember-client/src/lib.rs | 43 +- crates/ember-client/src/pipeline.rs | 355 +++++++++ tests/integration/Cargo.toml | 1 + tests/integration/src/client_typed_api.rs | 243 ++++++ tests/integration/src/main.rs | 1 + 9 files changed, 1663 insertions(+), 17 deletions(-) create mode 100644 crates/ember-client/README.md create mode 100644 crates/ember-client/src/commands.rs create mode 100644 crates/ember-client/src/pipeline.rs create mode 100644 tests/integration/src/client_typed_api.rs diff --git a/Cargo.lock b/Cargo.lock index 4abd497a..0441fa82 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -886,6 +886,7 @@ name = "ember-integration-tests" version = "0.4.8" dependencies = [ "bytes", + "ember-client", "ember-protocol", "ember-server", "prost-reflect", diff --git a/crates/ember-client/README.md b/crates/ember-client/README.md new file mode 100644 index 00000000..59e5374d --- /dev/null +++ b/crates/ember-client/README.md @@ -0,0 +1,96 @@ +# ember-client + +async Rust client for [ember](https://github.com/yourorg/ember), a high-performance distributed cache. + +## features + +- typed API for all common commands (strings, lists, hashes, sets, sorted sets) +- pipelining — batch commands into a single round-trip +- optional TLS via the `tls` feature (enabled by default) +- binary-safe values: inputs accept `&str`, `String`, `Vec`, or `bytes::Bytes` +- outputs return `Bytes` — zero-copy where possible + +## quick start + +```toml +[dependencies] +ember-client = "0.1" +tokio = { version = "1", features = ["full"] } +``` + +```rust +use ember_client::Client; + +#[tokio::main] +async fn main() -> Result<(), ember_client::ClientError> { + let mut client = Client::connect("127.0.0.1", 6379).await?; + + client.set("greeting", "hello").await?; + if let Some(value) = client.get("greeting").await? { + println!("{}", String::from_utf8_lossy(&value)); + } + + Ok(()) +} +``` + +## pipelining + +Send multiple commands in a single network round-trip: + +```rust +use ember_client::{Client, Pipeline}; + +#[tokio::main] +async fn main() -> Result<(), ember_client::ClientError> { + let mut client = Client::connect("127.0.0.1", 6379).await?; + + let frames = client.execute_pipeline( + Pipeline::new() + .set("a", "1") + .set("b", "2") + .get("a") + .get("b"), + ).await?; + + println!("{} responses", frames.len()); + Ok(()) +} +``` + +## commands + +| group | commands | +|-------------|----------| +| strings | `get`, `set`, `set_ex`, `del`, `exists`, `expire`, `persist`, `ttl`, `pttl`, `incr`, `decr`, `incrby`, `decrby`, `append`, `mget`, `mset`, `getdel` | +| lists | `lpush`, `rpush`, `lpop`, `rpop`, `lrange`, `llen` | +| hashes | `hset`, `hget`, `hgetall`, `hdel`, `hexists`, `hlen`, `hkeys`, `hvals` | +| sets | `sadd`, `srem`, `smembers`, `sismember`, `scard` | +| sorted sets | `zadd`, `zrange`, `zrange_withscores`, `zscore`, `zrank`, `zrem`, `zcard` | +| server | `ping`, `dbsize`, `flushdb` | +| raw | `send` — pass any command as `&[&str]` | + +## errors + +`ClientError` covers five cases: + +| variant | meaning | +|---------------|---------| +| `Io` | TCP-level failure | +| `Protocol` | unexpected RESP3 frame shape | +| `Server` | server returned an error reply (`WRONGTYPE`, `NOAUTH`, etc.) | +| `Disconnected`| server closed the connection | +| `Timeout` | connect or read timed out (5 s / 10 s defaults) | + +## tls + +```toml +ember-client = { version = "0.1", features = ["tls"] } +``` + +```rust +use ember_client::{Client, tls::TlsClientConfig}; + +let tls = TlsClientConfig::default(); // uses native root certs +let mut client = Client::connect_tls("my-ember-host", 6380, &tls).await?; +``` diff --git a/crates/ember-client/src/commands.rs b/crates/ember-client/src/commands.rs new file mode 100644 index 00000000..5761f1e2 --- /dev/null +++ b/crates/ember-client/src/commands.rs @@ -0,0 +1,896 @@ +//! Typed command API for the ember client. +//! +//! Provides ergonomic methods on [`Client`] for all common commands. Each +//! method takes strongly-typed inputs and returns decoded Rust values — no +//! manual frame inspection needed. +//! +//! Value inputs use `impl AsRef<[u8]>`, so `&str`, `String`, `Vec`, and +//! `bytes::Bytes` all work without explicit conversion. +//! +//! Value outputs use `Bytes` (binary-safe, reference-counted). Convert to a +//! `String` with `String::from_utf8(bytes.to_vec())` when you know the value +//! is UTF-8. + +use bytes::Bytes; +use ember_protocol::types::Frame; + +use crate::connection::{Client, ClientError}; +use crate::pipeline::Pipeline; + +// --- frame construction helpers --- + +fn cmd2(name: &'static [u8], a: &[u8]) -> Frame { + Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(name)), + Frame::Bulk(Bytes::copy_from_slice(a)), + ]) +} + +fn cmd3(name: &'static [u8], a: &[u8], b: &[u8]) -> Frame { + Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(name)), + Frame::Bulk(Bytes::copy_from_slice(a)), + Frame::Bulk(Bytes::copy_from_slice(b)), + ]) +} + +fn cmd4(name: &'static [u8], a: &[u8], b: &[u8], c: &[u8]) -> Frame { + Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(name)), + Frame::Bulk(Bytes::copy_from_slice(a)), + Frame::Bulk(Bytes::copy_from_slice(b)), + Frame::Bulk(Bytes::copy_from_slice(c)), + ]) +} + +// --- frame decoding helpers --- +// +// Each decoder maps a raw Frame to a typed result. +// `Frame::Error` → `ClientError::Server` +// Unexpected variants → `ClientError::Protocol` + +/// Decodes a bulk string response, returning `None` for null. +/// +/// Used by: GET, LPOP, RPOP, HGET, GETDEL. +fn optional_bytes(frame: Frame) -> Result, ClientError> { + match frame { + Frame::Bulk(b) => Ok(Some(b)), + Frame::Null => Ok(None), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected bulk or null, got {other:?}" + ))), + } +} + +/// Decodes an integer response. +/// +/// Used by: DEL, EXISTS, INCR, DECR, INCRBY, DECRBY, APPEND, LPUSH, RPUSH, +/// LLEN, HSET, HDEL, HLEN, SADD, SREM, SCARD, ZREM, ZCARD, ZADD, DBSIZE. +fn integer(frame: Frame) -> Result { + match frame { + Frame::Integer(n) => Ok(n), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected integer, got {other:?}" + ))), + } +} + +/// Decodes an integer response, returning `None` for null. +/// +/// Used by: ZRANK (returns null when member is absent). +fn optional_integer(frame: Frame) -> Result, ClientError> { + match frame { + Frame::Integer(n) => Ok(Some(n)), + Frame::Null => Ok(None), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected integer or null, got {other:?}" + ))), + } +} + +/// Decodes an integer 0/1 as a boolean. +/// +/// Used by: EXPIRE, PERSIST, HEXISTS, SISMEMBER. +fn bool_flag(frame: Frame) -> Result { + match frame { + Frame::Integer(1) => Ok(true), + Frame::Integer(0) => Ok(false), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected 0 or 1, got {other:?}" + ))), + } +} + +/// Decodes an OK simple string response. +/// +/// Used by: SET, MSET, FLUSHDB. +fn ok(frame: Frame) -> Result<(), ClientError> { + match frame { + Frame::Simple(s) if s == "OK" => Ok(()), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected +OK, got {other:?}" + ))), + } +} + +/// Decodes an array of bulk strings. +/// +/// Used by: LRANGE, SMEMBERS, HKEYS, HVALS, ZRANGE. +fn bytes_vec(frame: Frame) -> Result, ClientError> { + match frame { + Frame::Array(elems) => elems + .into_iter() + .map(|e| match e { + Frame::Bulk(b) => Ok(b), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected bulk in array, got {other:?}" + ))), + }) + .collect(), + Frame::Null => Ok(Vec::new()), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected array, got {other:?}" + ))), + } +} + +/// Decodes an array where individual elements may be null. +/// +/// Used by: MGET (missing keys come back as null within the array). +fn optional_bytes_vec(frame: Frame) -> Result>, ClientError> { + match frame { + Frame::Array(elems) => elems + .into_iter() + .map(|e| match e { + Frame::Bulk(b) => Ok(Some(b)), + Frame::Null => Ok(None), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "unexpected element in array: {other:?}" + ))), + }) + .collect(), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected array, got {other:?}" + ))), + } +} + +/// Decodes a flat array of alternating keys and values into pairs. +/// +/// Used by: HGETALL (field, value, field, value, ...). +fn pairs(frame: Frame) -> Result, ClientError> { + let elems = match frame { + Frame::Array(e) => e, + Frame::Null => return Ok(Vec::new()), + Frame::Error(e) => return Err(ClientError::Server(e)), + other => { + return Err(ClientError::Protocol(format!( + "expected array for pairs, got {other:?}" + ))) + } + }; + + if elems.len() % 2 != 0 { + return Err(ClientError::Protocol(format!( + "pairs array has odd length ({})", + elems.len() + ))); + } + + let mut result = Vec::with_capacity(elems.len() / 2); + let mut iter = elems.into_iter(); + while let (Some(k), Some(v)) = (iter.next(), iter.next()) { + let key = match k { + Frame::Bulk(b) => b, + other => { + return Err(ClientError::Protocol(format!( + "expected bulk key in pairs, got {other:?}" + ))) + } + }; + let val = match v { + Frame::Bulk(b) => b, + other => { + return Err(ClientError::Protocol(format!( + "expected bulk value in pairs, got {other:?}" + ))) + } + }; + result.push((key, val)); + } + Ok(result) +} + +/// Decodes a score returned as a bulk string float, returning `None` for null. +/// +/// Used by: ZSCORE. +fn optional_score(frame: Frame) -> Result, ClientError> { + match frame { + Frame::Bulk(b) => { + let s = std::str::from_utf8(&b) + .map_err(|_| ClientError::Protocol("score is not valid UTF-8".into()))?; + s.parse::() + .map(Some) + .map_err(|_| ClientError::Protocol(format!("score is not a valid float: {s:?}"))) + } + Frame::Null => Ok(None), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected bulk or null for score, got {other:?}" + ))), + } +} + +/// Decodes `ZRANGE WITHSCORES` — alternating member / score pairs. +/// +/// Used by: zrange_withscores. +fn scored_members(frame: Frame) -> Result, ClientError> { + let elems = match frame { + Frame::Array(e) => e, + Frame::Null => return Ok(Vec::new()), + Frame::Error(e) => return Err(ClientError::Server(e)), + other => { + return Err(ClientError::Protocol(format!( + "expected array for scored members, got {other:?}" + ))) + } + }; + + if elems.len() % 2 != 0 { + return Err(ClientError::Protocol(format!( + "WITHSCORES array has odd length ({})", + elems.len() + ))); + } + + let mut result = Vec::with_capacity(elems.len() / 2); + let mut iter = elems.into_iter(); + while let (Some(member_frame), Some(score_frame)) = (iter.next(), iter.next()) { + let member = match member_frame { + Frame::Bulk(b) => b, + other => { + return Err(ClientError::Protocol(format!( + "expected bulk member, got {other:?}" + ))) + } + }; + let score = match score_frame { + Frame::Bulk(b) => { + let s = std::str::from_utf8(&b) + .map_err(|_| ClientError::Protocol("score is not valid UTF-8".into()))?; + s.parse::().map_err(|_| { + ClientError::Protocol(format!("score is not a valid float: {s:?}")) + })? + } + other => { + return Err(ClientError::Protocol(format!( + "expected bulk score, got {other:?}" + ))) + } + }; + result.push((member, score)); + } + Ok(result) +} + +// --- typed command API --- + +impl Client { + // --- string commands --- + + /// Returns the value for `key`, or `None` if it does not exist. + pub async fn get(&mut self, key: &str) -> Result, ClientError> { + let frame = self.send_frame(cmd2(b"GET", key.as_bytes())).await?; + optional_bytes(frame) + } + + /// Sets `key` to `value` with no expiry. + pub async fn set(&mut self, key: &str, value: impl AsRef<[u8]>) -> Result<(), ClientError> { + let frame = self + .send_frame(cmd3(b"SET", key.as_bytes(), value.as_ref())) + .await?; + ok(frame) + } + + /// Sets `key` to `value` with an expiry of `seconds`. + pub async fn set_ex( + &mut self, + key: &str, + value: impl AsRef<[u8]>, + seconds: u64, + ) -> Result<(), ClientError> { + let secs = seconds.to_string(); + let frame = self + .send_frame(cmd4( + b"SETEX", + key.as_bytes(), + secs.as_bytes(), + value.as_ref(), + )) + .await?; + ok(frame) + } + + /// Deletes one or more keys. Returns the number of keys that were removed. + pub async fn del(&mut self, keys: &[&str]) -> Result { + let mut parts = Vec::with_capacity(1 + keys.len()); + parts.push(Frame::Bulk(Bytes::from_static(b"DEL"))); + for k in keys { + parts.push(Frame::Bulk(Bytes::copy_from_slice(k.as_bytes()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + integer(frame) + } + + /// Returns the number of supplied keys that exist. + pub async fn exists(&mut self, keys: &[&str]) -> Result { + let mut parts = Vec::with_capacity(1 + keys.len()); + parts.push(Frame::Bulk(Bytes::from_static(b"EXISTS"))); + for k in keys { + parts.push(Frame::Bulk(Bytes::copy_from_slice(k.as_bytes()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + integer(frame) + } + + /// Sets a timeout of `seconds` on `key`. Returns `true` if the timeout + /// was set, `false` if the key does not exist. + pub async fn expire(&mut self, key: &str, seconds: u64) -> Result { + let secs = seconds.to_string(); + let frame = self + .send_frame(cmd3(b"EXPIRE", key.as_bytes(), secs.as_bytes())) + .await?; + bool_flag(frame) + } + + /// Removes any existing timeout on `key`. Returns `true` if the timeout + /// was removed, `false` if the key has no timeout or does not exist. + pub async fn persist(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"PERSIST", key.as_bytes())).await?; + bool_flag(frame) + } + + /// Returns the remaining TTL in seconds. Returns `-2` if the key does not + /// exist, `-1` if it has no expiry. + pub async fn ttl(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"TTL", key.as_bytes())).await?; + integer(frame) + } + + /// Returns the remaining TTL in milliseconds. + pub async fn pttl(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"PTTL", key.as_bytes())).await?; + integer(frame) + } + + /// Increments the integer stored at `key` by 1. Returns the new value. + pub async fn incr(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"INCR", key.as_bytes())).await?; + integer(frame) + } + + /// Decrements the integer stored at `key` by 1. Returns the new value. + pub async fn decr(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"DECR", key.as_bytes())).await?; + integer(frame) + } + + /// Increments the integer stored at `key` by `delta`. Returns the new value. + pub async fn incrby(&mut self, key: &str, delta: i64) -> Result { + let d = delta.to_string(); + let frame = self + .send_frame(cmd3(b"INCRBY", key.as_bytes(), d.as_bytes())) + .await?; + integer(frame) + } + + /// Decrements the integer stored at `key` by `delta`. Returns the new value. + pub async fn decrby(&mut self, key: &str, delta: i64) -> Result { + let d = delta.to_string(); + let frame = self + .send_frame(cmd3(b"DECRBY", key.as_bytes(), d.as_bytes())) + .await?; + integer(frame) + } + + /// Appends `value` to the string at `key`. Returns the new length. + pub async fn append(&mut self, key: &str, value: impl AsRef<[u8]>) -> Result { + let frame = self + .send_frame(cmd3(b"APPEND", key.as_bytes(), value.as_ref())) + .await?; + integer(frame) + } + + /// Returns the values for multiple keys. Missing keys are `None`. + pub async fn mget(&mut self, keys: &[&str]) -> Result>, ClientError> { + let mut parts = Vec::with_capacity(1 + keys.len()); + parts.push(Frame::Bulk(Bytes::from_static(b"MGET"))); + for k in keys { + parts.push(Frame::Bulk(Bytes::copy_from_slice(k.as_bytes()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + optional_bytes_vec(frame) + } + + /// Sets multiple key-value pairs atomically. + pub async fn mset>(&mut self, pairs: &[(&str, V)]) -> Result<(), ClientError> { + let mut parts = Vec::with_capacity(1 + pairs.len() * 2); + parts.push(Frame::Bulk(Bytes::from_static(b"MSET"))); + for (k, v) in pairs { + parts.push(Frame::Bulk(Bytes::copy_from_slice(k.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(v.as_ref()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + ok(frame) + } + + /// Returns the value of `key` and deletes it atomically. Returns `None` + /// if the key does not exist. + pub async fn getdel(&mut self, key: &str) -> Result, ClientError> { + let frame = self.send_frame(cmd2(b"GETDEL", key.as_bytes())).await?; + optional_bytes(frame) + } + + // --- list commands --- + + /// Prepends `values` to the list at `key`. Returns the new list length. + pub async fn lpush>( + &mut self, + key: &str, + values: &[V], + ) -> Result { + let frame = self + .send_frame(cmd_key_values(b"LPUSH", key, values)) + .await?; + integer(frame) + } + + /// Appends `values` to the list at `key`. Returns the new list length. + pub async fn rpush>( + &mut self, + key: &str, + values: &[V], + ) -> Result { + let frame = self + .send_frame(cmd_key_values(b"RPUSH", key, values)) + .await?; + integer(frame) + } + + /// Removes and returns the first element of the list at `key`. + pub async fn lpop(&mut self, key: &str) -> Result, ClientError> { + let frame = self.send_frame(cmd2(b"LPOP", key.as_bytes())).await?; + optional_bytes(frame) + } + + /// Removes and returns the last element of the list at `key`. + pub async fn rpop(&mut self, key: &str) -> Result, ClientError> { + let frame = self.send_frame(cmd2(b"RPOP", key.as_bytes())).await?; + optional_bytes(frame) + } + + /// Returns the elements of the list between `start` and `stop` (inclusive). + pub async fn lrange( + &mut self, + key: &str, + start: i64, + stop: i64, + ) -> Result, ClientError> { + let s = start.to_string(); + let e = stop.to_string(); + let frame = self + .send_frame(cmd4(b"LRANGE", key.as_bytes(), s.as_bytes(), e.as_bytes())) + .await?; + bytes_vec(frame) + } + + /// Returns the length of the list at `key`. Returns 0 for missing keys. + pub async fn llen(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"LLEN", key.as_bytes())).await?; + integer(frame) + } + + // --- hash commands --- + + /// Sets `field`/`value` pairs on the hash at `key`. Returns the number of + /// new fields added. + pub async fn hset>( + &mut self, + key: &str, + pairs: &[(&str, V)], + ) -> Result { + let mut parts = Vec::with_capacity(2 + pairs.len() * 2); + parts.push(Frame::Bulk(Bytes::from_static(b"HSET"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + for (field, val) in pairs { + parts.push(Frame::Bulk(Bytes::copy_from_slice(field.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(val.as_ref()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + integer(frame) + } + + /// Returns the value at `field` in the hash at `key`, or `None` if either + /// does not exist. + pub async fn hget(&mut self, key: &str, field: &str) -> Result, ClientError> { + let frame = self + .send_frame(cmd3(b"HGET", key.as_bytes(), field.as_bytes())) + .await?; + optional_bytes(frame) + } + + /// Returns all field-value pairs in the hash at `key`. + pub async fn hgetall(&mut self, key: &str) -> Result, ClientError> { + let frame = self.send_frame(cmd2(b"HGETALL", key.as_bytes())).await?; + pairs(frame) + } + + /// Deletes `fields` from the hash at `key`. Returns the number removed. + pub async fn hdel(&mut self, key: &str, fields: &[&str]) -> Result { + let frame = self + .send_frame(cmd_key_and_keys(b"HDEL", key, fields)) + .await?; + integer(frame) + } + + /// Returns `true` if `field` exists in the hash at `key`. + pub async fn hexists(&mut self, key: &str, field: &str) -> Result { + let frame = self + .send_frame(cmd3(b"HEXISTS", key.as_bytes(), field.as_bytes())) + .await?; + bool_flag(frame) + } + + /// Returns the number of fields in the hash at `key`. + pub async fn hlen(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"HLEN", key.as_bytes())).await?; + integer(frame) + } + + /// Returns all field names in the hash at `key`. + pub async fn hkeys(&mut self, key: &str) -> Result, ClientError> { + let frame = self.send_frame(cmd2(b"HKEYS", key.as_bytes())).await?; + bytes_vec(frame) + } + + /// Returns all values in the hash at `key`. + pub async fn hvals(&mut self, key: &str) -> Result, ClientError> { + let frame = self.send_frame(cmd2(b"HVALS", key.as_bytes())).await?; + bytes_vec(frame) + } + + // --- set commands --- + + /// Adds `members` to the set at `key`. Returns the number added. + pub async fn sadd(&mut self, key: &str, members: &[&str]) -> Result { + let frame = self + .send_frame(cmd_key_and_keys(b"SADD", key, members)) + .await?; + integer(frame) + } + + /// Removes `members` from the set at `key`. Returns the number removed. + pub async fn srem(&mut self, key: &str, members: &[&str]) -> Result { + let frame = self + .send_frame(cmd_key_and_keys(b"SREM", key, members)) + .await?; + integer(frame) + } + + /// Returns all members of the set at `key`. + pub async fn smembers(&mut self, key: &str) -> Result, ClientError> { + let frame = self.send_frame(cmd2(b"SMEMBERS", key.as_bytes())).await?; + bytes_vec(frame) + } + + /// Returns `true` if `member` belongs to the set at `key`. + pub async fn sismember(&mut self, key: &str, member: &str) -> Result { + let frame = self + .send_frame(cmd3(b"SISMEMBER", key.as_bytes(), member.as_bytes())) + .await?; + bool_flag(frame) + } + + /// Returns the cardinality (number of members) of the set at `key`. + pub async fn scard(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"SCARD", key.as_bytes())).await?; + integer(frame) + } + + // --- sorted set commands --- + + /// Adds `members` to the sorted set at `key`. Each member is a + /// `(score, name)` pair. Returns the number of new members added. + pub async fn zadd(&mut self, key: &str, members: &[(f64, &str)]) -> Result { + let mut parts = Vec::with_capacity(2 + members.len() * 2); + parts.push(Frame::Bulk(Bytes::from_static(b"ZADD"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + for (score, member) in members { + let s = score.to_string(); + parts.push(Frame::Bulk(Bytes::copy_from_slice(s.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(member.as_bytes()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + integer(frame) + } + + /// Returns members of the sorted set between rank `start` and `stop` + /// (0-based, inclusive). + pub async fn zrange( + &mut self, + key: &str, + start: i64, + stop: i64, + ) -> Result, ClientError> { + let s = start.to_string(); + let e = stop.to_string(); + let frame = self + .send_frame(cmd4(b"ZRANGE", key.as_bytes(), s.as_bytes(), e.as_bytes())) + .await?; + bytes_vec(frame) + } + + /// Like `zrange`, but also returns scores as `(member, score)` pairs. + pub async fn zrange_withscores( + &mut self, + key: &str, + start: i64, + stop: i64, + ) -> Result, ClientError> { + let s = start.to_string(); + let e = stop.to_string(); + let frame = self + .send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"ZRANGE")), + Frame::Bulk(Bytes::copy_from_slice(key.as_bytes())), + Frame::Bulk(Bytes::copy_from_slice(s.as_bytes())), + Frame::Bulk(Bytes::copy_from_slice(e.as_bytes())), + Frame::Bulk(Bytes::from_static(b"WITHSCORES")), + ])) + .await?; + scored_members(frame) + } + + /// Returns the score of `member` in the sorted set, or `None` if absent. + pub async fn zscore(&mut self, key: &str, member: &str) -> Result, ClientError> { + let frame = self + .send_frame(cmd3(b"ZSCORE", key.as_bytes(), member.as_bytes())) + .await?; + optional_score(frame) + } + + /// Returns the rank of `member` in the sorted set (0-based, ascending + /// score order), or `None` if absent. + pub async fn zrank(&mut self, key: &str, member: &str) -> Result, ClientError> { + let frame = self + .send_frame(cmd3(b"ZRANK", key.as_bytes(), member.as_bytes())) + .await?; + optional_integer(frame) + } + + /// Removes `members` from the sorted set at `key`. Returns the number + /// removed. + pub async fn zrem(&mut self, key: &str, members: &[&str]) -> Result { + let frame = self + .send_frame(cmd_key_and_keys(b"ZREM", key, members)) + .await?; + integer(frame) + } + + /// Returns the number of members in the sorted set at `key`. + pub async fn zcard(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"ZCARD", key.as_bytes())).await?; + integer(frame) + } + + // --- server commands --- + + /// Sends a `PING` and expects a `PONG` response. + pub async fn ping(&mut self) -> Result<(), ClientError> { + let frame = self + .send_frame(Frame::Array(vec![Frame::Bulk(Bytes::from_static(b"PING"))])) + .await?; + match frame { + Frame::Simple(s) if s == "PONG" => Ok(()), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected PONG, got {other:?}" + ))), + } + } + + /// Returns the number of keys in the current database. + pub async fn dbsize(&mut self) -> Result { + let frame = self + .send_frame(Frame::Array(vec![Frame::Bulk(Bytes::from_static( + b"DBSIZE", + ))])) + .await?; + integer(frame) + } + + /// Removes all keys from the current database. + pub async fn flushdb(&mut self) -> Result<(), ClientError> { + let frame = self + .send_frame(Frame::Array(vec![Frame::Bulk(Bytes::from_static( + b"FLUSHDB", + ))])) + .await?; + ok(frame) + } + + // --- pipeline --- + + /// Executes all commands queued in `pipeline` as a single batch. + /// + /// Returns one raw [`Frame`] per command in the same order they were + /// queued. An empty pipeline returns an empty `Vec` without touching the + /// network. + /// + /// Use this when you need full control over decoding the responses (e.g. + /// mixed command types in one batch). For homogeneous batches the typed + /// builder methods on [`Pipeline`] pair well with a manual decode loop. + pub async fn execute_pipeline( + &mut self, + pipeline: Pipeline, + ) -> Result, ClientError> { + if pipeline.is_empty() { + return Ok(Vec::new()); + } + self.send_batch(&pipeline.cmds).await + } +} + +// --- shared frame construction helpers used by multiple command groups --- + +fn cmd_key_and_keys(cmd: &'static [u8], key: &str, rest: &[&str]) -> Frame { + let mut parts = Vec::with_capacity(2 + rest.len()); + parts.push(Frame::Bulk(Bytes::from_static(cmd))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + for s in rest { + parts.push(Frame::Bulk(Bytes::copy_from_slice(s.as_bytes()))); + } + Frame::Array(parts) +} + +fn cmd_key_values>(cmd: &'static [u8], key: &str, values: &[V]) -> Frame { + let mut parts = Vec::with_capacity(2 + values.len()); + parts.push(Frame::Bulk(Bytes::from_static(cmd))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + for v in values { + parts.push(Frame::Bulk(Bytes::copy_from_slice(v.as_ref()))); + } + Frame::Array(parts) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn bulk(b: &[u8]) -> Frame { + Frame::Bulk(Bytes::copy_from_slice(b)) + } + + // --- optional_bytes --- + + #[test] + fn optional_bytes_bulk() { + let f = Frame::Bulk(Bytes::from_static(b"hello")); + assert_eq!( + optional_bytes(f).unwrap(), + Some(Bytes::from_static(b"hello")) + ); + } + + #[test] + fn optional_bytes_null() { + assert_eq!(optional_bytes(Frame::Null).unwrap(), None); + } + + #[test] + fn optional_bytes_server_error() { + let f = Frame::Error("WRONGTYPE".into()); + assert!(matches!(optional_bytes(f), Err(ClientError::Server(_)))); + } + + #[test] + fn optional_bytes_unexpected() { + let f = Frame::Integer(1); + assert!(matches!(optional_bytes(f), Err(ClientError::Protocol(_)))); + } + + // --- bool_flag --- + + #[test] + fn bool_flag_one() { + assert_eq!(bool_flag(Frame::Integer(1)).unwrap(), true); + } + + #[test] + fn bool_flag_zero() { + assert_eq!(bool_flag(Frame::Integer(0)).unwrap(), false); + } + + #[test] + fn bool_flag_unexpected_value() { + let f = Frame::Integer(42); + assert!(matches!(bool_flag(f), Err(ClientError::Protocol(_)))); + } + + // --- pairs --- + + #[test] + fn pairs_even_array() { + let f = Frame::Array(vec![ + bulk(b"field1"), + bulk(b"val1"), + bulk(b"field2"), + bulk(b"val2"), + ]); + let result = pairs(f).unwrap(); + assert_eq!(result.len(), 2); + assert_eq!(result[0].0, Bytes::from_static(b"field1")); + assert_eq!(result[0].1, Bytes::from_static(b"val1")); + } + + #[test] + fn pairs_odd_array_is_error() { + let f = Frame::Array(vec![bulk(b"orphan")]); + assert!(matches!(pairs(f), Err(ClientError::Protocol(_)))); + } + + #[test] + fn pairs_empty_array() { + assert_eq!(pairs(Frame::Array(vec![])).unwrap(), vec![]); + } + + #[test] + fn pairs_null() { + assert_eq!(pairs(Frame::Null).unwrap(), vec![]); + } + + // --- optional_score --- + + #[test] + fn optional_score_valid_float() { + let f = Frame::Bulk(Bytes::from_static(b"3.14")); + let s = optional_score(f).unwrap(); + assert!((s.unwrap() - 3.14).abs() < f64::EPSILON); + } + + #[test] + fn optional_score_null() { + assert_eq!(optional_score(Frame::Null).unwrap(), None); + } + + #[test] + fn optional_score_malformed() { + let f = Frame::Bulk(Bytes::from_static(b"notanumber")); + assert!(matches!(optional_score(f), Err(ClientError::Protocol(_)))); + } + + // --- optional_bytes_vec (MGET) --- + + #[test] + fn optional_bytes_vec_null_elements() { + let f = Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"value")), + Frame::Null, + Frame::Bulk(Bytes::from_static(b"another")), + ]); + let result = optional_bytes_vec(f).unwrap(); + assert_eq!(result.len(), 3); + assert!(result[0].is_some()); + assert!(result[1].is_none()); + assert!(result[2].is_some()); + } +} diff --git a/crates/ember-client/src/connection.rs b/crates/ember-client/src/connection.rs index 45e66453..d01a3c67 100644 --- a/crates/ember-client/src/connection.rs +++ b/crates/ember-client/src/connection.rs @@ -36,6 +36,9 @@ pub enum ClientError { #[error("protocol error: {0}")] Protocol(String), + #[error("server error: {0}")] + Server(String), + #[error("server disconnected")] Disconnected, @@ -169,20 +172,44 @@ impl Client { /// # } /// ``` pub async fn send(&mut self, args: &[&str]) -> Result { - let parts: Vec = args + let parts = args .iter() - .map(|t| Frame::Bulk(bytes::Bytes::from(t.to_string()))) + .map(|t| Frame::Bulk(bytes::Bytes::copy_from_slice(t.as_bytes()))) .collect(); - let frame = Frame::Array(parts); + self.send_frame(Frame::Array(parts)).await + } + /// Sends a single pre-built frame and returns the response. + /// + /// Used internally by typed command methods to avoid re-encoding + /// arguments through `&[&str]`. + pub(crate) async fn send_frame(&mut self, frame: Frame) -> Result { self.write_buf.clear(); frame.serialize(&mut self.write_buf); self.transport.write_all(&self.write_buf).await?; self.transport.flush().await?; - self.read_response().await } + /// Writes all frames in a single flush, then reads one response per frame. + /// + /// This is the core of pipelining: batching multiple commands into one + /// syscall and reading responses sequentially afterwards. + pub(crate) async fn send_batch(&mut self, frames: &[Frame]) -> Result, ClientError> { + self.write_buf.clear(); + for frame in frames { + frame.serialize(&mut self.write_buf); + } + self.transport.write_all(&self.write_buf).await?; + self.transport.flush().await?; + + let mut results = Vec::with_capacity(frames.len()); + for _ in 0..frames.len() { + results.push(self.read_response().await?); + } + Ok(results) + } + /// Authenticates with the server using the `AUTH` command. /// /// Returns `Ok(())` on success, or `ClientError::AuthFailed` if the server @@ -190,15 +217,10 @@ impl Client { pub async fn auth(&mut self, password: &str) -> Result<(), ClientError> { let frame = Frame::Array(vec![ Frame::Bulk(bytes::Bytes::from_static(b"AUTH")), - Frame::Bulk(bytes::Bytes::from(password.to_string())), + Frame::Bulk(bytes::Bytes::copy_from_slice(password.as_bytes())), ]); - self.write_buf.clear(); - frame.serialize(&mut self.write_buf); - self.transport.write_all(&self.write_buf).await?; - self.transport.flush().await?; - - match self.read_response().await? { + match self.send_frame(frame).await? { Frame::Simple(s) if s == "OK" => Ok(()), Frame::Error(e) => Err(ClientError::AuthFailed(e)), _ => Err(ClientError::AuthFailed( diff --git a/crates/ember-client/src/lib.rs b/crates/ember-client/src/lib.rs index 2b9d32ed..a729d466 100644 --- a/crates/ember-client/src/lib.rs +++ b/crates/ember-client/src/lib.rs @@ -1,9 +1,10 @@ -//! ember-client: async RESP3 client for ember. +//! ember-client: async Rust client for ember. //! -//! Provides a simple async client for connecting to an ember server over -//! TCP (or TLS), sending commands as RESP3 frames, and reading responses. +//! Connects to an ember server over TCP (or TLS) and exposes a typed async +//! API covering all common commands. Raw RESP3 frames are available via +//! [`Client::send`] when you need something not covered by the typed methods. //! -//! # Example +//! # Quick start //! //! ```no_run //! use ember_client::Client; @@ -11,17 +12,47 @@ //! #[tokio::main] //! async fn main() -> Result<(), ember_client::ClientError> { //! let mut client = Client::connect("127.0.0.1", 6379).await?; -//! let response = client.send(&["PING"]).await?; -//! println!("{response:?}"); +//! +//! client.set("greeting", "hello").await?; +//! let value = client.get("greeting").await?; +//! println!("{value:?}"); // Some(b"hello") +//! +//! Ok(()) +//! } +//! ``` +//! +//! # Pipelining +//! +//! Send multiple commands in one round-trip using [`Pipeline`]: +//! +//! ```no_run +//! use ember_client::{Client, Pipeline}; +//! +//! #[tokio::main] +//! async fn main() -> Result<(), ember_client::ClientError> { +//! let mut client = Client::connect("127.0.0.1", 6379).await?; +//! +//! let frames = client.execute_pipeline( +//! Pipeline::new() +//! .set("a", "1") +//! .set("b", "2") +//! .get("a") +//! .get("b"), +//! ).await?; +//! +//! println!("{} responses", frames.len()); //! Ok(()) //! } //! ``` +mod commands; mod connection; +mod pipeline; #[cfg(feature = "tls")] pub mod tls; pub use connection::{Client, ClientError}; pub use ember_protocol::types::Frame; +pub use pipeline::Pipeline; #[cfg(feature = "tls")] pub use tls::TlsClientConfig; diff --git a/crates/ember-client/src/pipeline.rs b/crates/ember-client/src/pipeline.rs new file mode 100644 index 00000000..0346d40d --- /dev/null +++ b/crates/ember-client/src/pipeline.rs @@ -0,0 +1,355 @@ +//! Command pipeline builder. +//! +//! A [`Pipeline`] queues multiple commands and sends them in a single network +//! round-trip. This dramatically reduces latency when you need to issue many +//! independent commands in sequence. +//! +//! # Example +//! +//! ```no_run +//! use ember_client::{Client, Pipeline}; +//! +//! #[tokio::main] +//! async fn main() -> Result<(), ember_client::ClientError> { +//! let mut client = Client::connect("127.0.0.1", 6379).await?; +//! +//! let results = client.execute_pipeline( +//! Pipeline::new() +//! .set("greeting", "hello") +//! .get("greeting") +//! .incr("visits"), +//! ).await?; +//! +//! println!("{} responses", results.len()); +//! Ok(()) +//! } +//! ``` + +use bytes::Bytes; +use ember_protocol::types::Frame; + +/// A batch of commands to be sent in a single network round-trip. +/// +/// Build the pipeline with the typed builder methods, then execute it +/// with [`Client::execute_pipeline`]. +pub struct Pipeline { + pub(crate) cmds: Vec, +} + +impl Pipeline { + /// Creates an empty pipeline. + pub fn new() -> Self { + Self { cmds: Vec::new() } + } + + /// Returns the number of queued commands. + pub fn len(&self) -> usize { + self.cmds.len() + } + + /// Returns `true` if no commands have been queued. + pub fn is_empty(&self) -> bool { + self.cmds.is_empty() + } + + /// Queues a raw command from string slices. + /// + /// Useful for commands not covered by the typed builders. + pub fn send(self, args: &[&str]) -> Self { + let parts = args + .iter() + .map(|s| Frame::Bulk(Bytes::copy_from_slice(s.as_bytes()))) + .collect(); + self.push(Frame::Array(parts)) + } + + // --- string commands --- + + /// Queues a `GET key` command. + pub fn get(self, key: &str) -> Self { + self.push(array2(b"GET", key.as_bytes())) + } + + /// Queues a `SET key value` command. + pub fn set(self, key: &str, value: impl AsRef<[u8]>) -> Self { + self.push(array3(b"SET", key.as_bytes(), value.as_ref())) + } + + /// Queues a `DEL key [key ...]` command. + pub fn del(self, keys: &[&str]) -> Self { + self.push(array_with_keys(b"DEL", keys)) + } + + /// Queues an `EXPIRE key seconds` command. + pub fn expire(self, key: &str, seconds: u64) -> Self { + let secs = seconds.to_string(); + self.push(array3(b"EXPIRE", key.as_bytes(), secs.as_bytes())) + } + + /// Queues an `INCR key` command. + pub fn incr(self, key: &str) -> Self { + self.push(array2(b"INCR", key.as_bytes())) + } + + /// Queues a `DECR key` command. + pub fn decr(self, key: &str) -> Self { + self.push(array2(b"DECR", key.as_bytes())) + } + + /// Queues an `INCRBY key delta` command. + pub fn incrby(self, key: &str, delta: i64) -> Self { + let d = delta.to_string(); + self.push(array3(b"INCRBY", key.as_bytes(), d.as_bytes())) + } + + /// Queues a `PING` command. + pub fn ping(self) -> Self { + self.push(Frame::Array(vec![Frame::Bulk(Bytes::from_static(b"PING"))])) + } + + /// Queues an `EXISTS key [key ...]` command. + pub fn exists(self, keys: &[&str]) -> Self { + self.push(array_with_keys(b"EXISTS", keys)) + } + + /// Queues a `TTL key` command. + pub fn ttl(self, key: &str) -> Self { + self.push(array2(b"TTL", key.as_bytes())) + } + + // --- list commands --- + + /// Queues an `LPUSH key value [value ...]` command. + pub fn lpush>(self, key: &str, values: &[V]) -> Self { + self.push(array_with_cmd_key_values(b"LPUSH", key, values)) + } + + /// Queues an `RPUSH key value [value ...]` command. + pub fn rpush>(self, key: &str, values: &[V]) -> Self { + self.push(array_with_cmd_key_values(b"RPUSH", key, values)) + } + + /// Queues an `LPOP key` command. + pub fn lpop(self, key: &str) -> Self { + self.push(array2(b"LPOP", key.as_bytes())) + } + + /// Queues an `RPOP key` command. + pub fn rpop(self, key: &str) -> Self { + self.push(array2(b"RPOP", key.as_bytes())) + } + + /// Queues an `LLEN key` command. + pub fn llen(self, key: &str) -> Self { + self.push(array2(b"LLEN", key.as_bytes())) + } + + // --- hash commands --- + + /// Queues an `HGET key field` command. + pub fn hget(self, key: &str, field: &str) -> Self { + self.push(array3(b"HGET", key.as_bytes(), field.as_bytes())) + } + + /// Queues an `HSET key field value [field value ...]` command. + pub fn hset>(self, key: &str, pairs: &[(&str, V)]) -> Self { + let mut parts = Vec::with_capacity(2 + pairs.len() * 2); + parts.push(Frame::Bulk(Bytes::from_static(b"HSET"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + for (field, val) in pairs { + parts.push(Frame::Bulk(Bytes::copy_from_slice(field.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(val.as_ref()))); + } + self.push(Frame::Array(parts)) + } + + /// Queues an `HDEL key field [field ...]` command. + pub fn hdel(self, key: &str, fields: &[&str]) -> Self { + self.push(array_with_key_and_keys(b"HDEL", key, fields)) + } + + // --- set commands --- + + /// Queues an `SADD key member [member ...]` command. + pub fn sadd(self, key: &str, members: &[&str]) -> Self { + self.push(array_with_key_and_keys(b"SADD", key, members)) + } + + /// Queues an `SREM key member [member ...]` command. + pub fn srem(self, key: &str, members: &[&str]) -> Self { + self.push(array_with_key_and_keys(b"SREM", key, members)) + } + + /// Queues an `SCARD key` command. + pub fn scard(self, key: &str) -> Self { + self.push(array2(b"SCARD", key.as_bytes())) + } + + // --- sorted set commands --- + + /// Queues a `ZADD key score member [score member ...]` command. + pub fn zadd(self, key: &str, members: &[(f64, &str)]) -> Self { + let mut parts = Vec::with_capacity(2 + members.len() * 2); + parts.push(Frame::Bulk(Bytes::from_static(b"ZADD"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + for (score, member) in members { + let s = score.to_string(); + parts.push(Frame::Bulk(Bytes::copy_from_slice(s.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(member.as_bytes()))); + } + self.push(Frame::Array(parts)) + } + + /// Queues a `ZCARD key` command. + pub fn zcard(self, key: &str) -> Self { + self.push(array2(b"ZCARD", key.as_bytes())) + } + + /// Queues a `ZSCORE key member` command. + pub fn zscore(self, key: &str, member: &str) -> Self { + self.push(array3(b"ZSCORE", key.as_bytes(), member.as_bytes())) + } + + // --- internal --- + + fn push(mut self, frame: Frame) -> Self { + self.cmds.push(frame); + self + } +} + +impl Default for Pipeline { + fn default() -> Self { + Self::new() + } +} + +// --- frame construction helpers --- +// These exist so each builder method stays a one-liner. + +fn array2(cmd: &'static [u8], a: &[u8]) -> Frame { + Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(cmd)), + Frame::Bulk(Bytes::copy_from_slice(a)), + ]) +} + +fn array3(cmd: &'static [u8], a: &[u8], b: &[u8]) -> Frame { + Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(cmd)), + Frame::Bulk(Bytes::copy_from_slice(a)), + Frame::Bulk(Bytes::copy_from_slice(b)), + ]) +} + +/// `CMD key1 key2 ...` +fn array_with_keys(cmd: &'static [u8], keys: &[&str]) -> Frame { + let mut parts = Vec::with_capacity(1 + keys.len()); + parts.push(Frame::Bulk(Bytes::from_static(cmd))); + for k in keys { + parts.push(Frame::Bulk(Bytes::copy_from_slice(k.as_bytes()))); + } + Frame::Array(parts) +} + +/// `CMD key field1 field2 ...` +fn array_with_key_and_keys(cmd: &'static [u8], key: &str, rest: &[&str]) -> Frame { + let mut parts = Vec::with_capacity(2 + rest.len()); + parts.push(Frame::Bulk(Bytes::from_static(cmd))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + for s in rest { + parts.push(Frame::Bulk(Bytes::copy_from_slice(s.as_bytes()))); + } + Frame::Array(parts) +} + +/// `CMD key value1 value2 ...` +fn array_with_cmd_key_values>(cmd: &'static [u8], key: &str, values: &[V]) -> Frame { + let mut parts = Vec::with_capacity(2 + values.len()); + parts.push(Frame::Bulk(Bytes::from_static(cmd))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + for v in values { + parts.push(Frame::Bulk(Bytes::copy_from_slice(v.as_ref()))); + } + Frame::Array(parts) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn bulk(b: &[u8]) -> Frame { + Frame::Bulk(Bytes::copy_from_slice(b)) + } + + #[test] + fn get_produces_correct_frame() { + let pipe = Pipeline::new().get("mykey"); + assert_eq!(pipe.len(), 1); + assert_eq!( + pipe.cmds[0], + Frame::Array(vec![bulk(b"GET"), bulk(b"mykey")]) + ); + } + + #[test] + fn del_multiple_keys() { + let pipe = Pipeline::new().del(&["a", "b"]); + assert_eq!(pipe.len(), 1); + match &pipe.cmds[0] { + Frame::Array(parts) => { + assert_eq!(parts.len(), 3); // DEL + 2 keys + assert_eq!(parts[0], bulk(b"DEL")); + assert_eq!(parts[1], bulk(b"a")); + assert_eq!(parts[2], bulk(b"b")); + } + other => panic!("expected Array, got {other:?}"), + } + } + + #[test] + fn chaining_accumulates_commands() { + let pipe = Pipeline::new() + .ping() + .get("k1") + .set("k2", "v2") + .incr("counter"); + assert_eq!(pipe.len(), 4); + } + + #[test] + fn empty_pipeline() { + let pipe = Pipeline::new(); + assert!(pipe.is_empty()); + assert_eq!(pipe.len(), 0); + } + + #[test] + fn hset_pairs_layout() { + let pipe = Pipeline::new().hset("myhash", &[("field1", "val1"), ("field2", "val2")]); + match &pipe.cmds[0] { + Frame::Array(parts) => { + assert_eq!(parts.len(), 6); // HSET + key + 2*(field+val) + assert_eq!(parts[0], bulk(b"HSET")); + assert_eq!(parts[1], bulk(b"myhash")); + assert_eq!(parts[2], bulk(b"field1")); + assert_eq!(parts[3], bulk(b"val1")); + } + other => panic!("expected Array, got {other:?}"), + } + } + + #[test] + fn zadd_score_member_layout() { + let pipe = Pipeline::new().zadd("leaderboard", &[(1.5, "alice"), (2.0, "bob")]); + match &pipe.cmds[0] { + Frame::Array(parts) => { + assert_eq!(parts.len(), 6); // ZADD + key + 2*(score+member) + assert_eq!(parts[0], bulk(b"ZADD")); + assert_eq!(parts[2], bulk(b"1.5")); + assert_eq!(parts[3], bulk(b"alice")); + } + other => panic!("expected Array, got {other:?}"), + } + } +} diff --git a/tests/integration/Cargo.toml b/tests/integration/Cargo.toml index 7b286e85..62a6adf3 100644 --- a/tests/integration/Cargo.toml +++ b/tests/integration/Cargo.toml @@ -13,6 +13,7 @@ harness = true [dependencies] tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "io-util", "time"] } bytes = { workspace = true } +ember-client = { path = "../../crates/ember-client", default-features = false } ember-protocol = { workspace = true } ember-server = { path = "../../crates/ember-server", default-features = false } prost-reflect = { workspace = true } diff --git a/tests/integration/src/client_typed_api.rs b/tests/integration/src/client_typed_api.rs new file mode 100644 index 00000000..170ac479 --- /dev/null +++ b/tests/integration/src/client_typed_api.rs @@ -0,0 +1,243 @@ +//! Integration tests for the typed ember-client API. +//! +//! Each test spins up a real server subprocess via `TestServer::start()` and +//! connects with the typed `Client`. These tests exercise the full stack — +//! TCP, RESP3 framing, command dispatch, and response decoding. + +use ember_client::{Client, ClientError, Pipeline}; + +use crate::helpers::TestServer; + +// Helper: connect a typed client to a freshly started server. +async fn connect() -> (TestServer, Client) { + let server = TestServer::start(); + let client = Client::connect("127.0.0.1", server.port) + .await + .expect("connect failed"); + (server, client) +} + +// --- string commands --- + +#[tokio::test] +async fn get_set_roundtrip() { + let (_server, mut client) = connect().await; + client.set("greeting", "hello").await.unwrap(); + let val = client.get("greeting").await.unwrap(); + assert_eq!(val.as_deref(), Some(b"hello" as &[u8])); +} + +#[tokio::test] +async fn get_missing_key_returns_none() { + let (_server, mut client) = connect().await; + let val = client.get("no_such_key").await.unwrap(); + assert!(val.is_none()); +} + +#[tokio::test] +async fn del_removes_keys_and_returns_count() { + let (_server, mut client) = connect().await; + client.set("k1", "a").await.unwrap(); + client.set("k2", "b").await.unwrap(); + let removed = client.del(&["k1", "k2", "k3"]).await.unwrap(); + assert_eq!(removed, 2); // k3 did not exist + assert!(client.get("k1").await.unwrap().is_none()); +} + +#[tokio::test] +async fn incr_decr_sequencing() { + let (_server, mut client) = connect().await; + client.set("counter", "10").await.unwrap(); + assert_eq!(client.incr("counter").await.unwrap(), 11); + assert_eq!(client.incrby("counter", 4).await.unwrap(), 15); + assert_eq!(client.decr("counter").await.unwrap(), 14); + assert_eq!(client.decrby("counter", 4).await.unwrap(), 10); +} + +#[tokio::test] +async fn mget_with_missing_key() { + let (_server, mut client) = connect().await; + client.set("present", "yes").await.unwrap(); + let vals = client.mget(&["present", "absent"]).await.unwrap(); + assert_eq!(vals.len(), 2); + assert_eq!(vals[0].as_deref(), Some(b"yes" as &[u8])); + assert!(vals[1].is_none()); +} + +#[tokio::test] +async fn mset_then_mget() { + let (_server, mut client) = connect().await; + client + .mset(&[("x", "1"), ("y", "2"), ("z", "3")]) + .await + .unwrap(); + let vals = client.mget(&["x", "y", "z"]).await.unwrap(); + assert_eq!(vals[0].as_deref(), Some(b"1" as &[u8])); + assert_eq!(vals[1].as_deref(), Some(b"2" as &[u8])); + assert_eq!(vals[2].as_deref(), Some(b"3" as &[u8])); +} + +#[tokio::test] +async fn getdel_removes_key() { + let (_server, mut client) = connect().await; + client.set("temp", "value").await.unwrap(); + let got = client.getdel("temp").await.unwrap(); + assert_eq!(got.as_deref(), Some(b"value" as &[u8])); + assert!(client.get("temp").await.unwrap().is_none()); +} + +#[tokio::test] +async fn expire_and_ttl() { + let (_server, mut client) = connect().await; + client.set("expiring", "soon").await.unwrap(); + let set = client.expire("expiring", 60).await.unwrap(); + assert!(set); + let ttl = client.ttl("expiring").await.unwrap(); + assert!(ttl > 0 && ttl <= 60); +} + +// --- list commands --- + +#[tokio::test] +async fn lpush_rpush_lrange() { + let (_server, mut client) = connect().await; + client.rpush("mylist", &["a", "b", "c"]).await.unwrap(); + let items = client.lrange("mylist", 0, -1).await.unwrap(); + assert_eq!(items.len(), 3); + assert_eq!(items[0].as_ref(), b"a"); + assert_eq!(items[2].as_ref(), b"c"); +} + +#[tokio::test] +async fn lpop_rpop() { + let (_server, mut client) = connect().await; + client.rpush("q", &["first", "second"]).await.unwrap(); + let head = client.lpop("q").await.unwrap(); + let tail = client.rpop("q").await.unwrap(); + assert_eq!(head.as_deref(), Some(b"first" as &[u8])); + assert_eq!(tail.as_deref(), Some(b"second" as &[u8])); + assert!(client.lpop("q").await.unwrap().is_none()); +} + +// --- hash commands --- + +#[tokio::test] +async fn hset_hgetall_roundtrip() { + let (_server, mut client) = connect().await; + let added = client + .hset("user:1", &[("name", "alice"), ("age", "30")]) + .await + .unwrap(); + assert_eq!(added, 2); + + let mut all = client.hgetall("user:1").await.unwrap(); + // order is not guaranteed — sort by field name for comparison + all.sort_by(|a, b| a.0.cmp(&b.0)); + assert_eq!(all.len(), 2); + assert_eq!(all[0].0.as_ref(), b"age"); + assert_eq!(all[0].1.as_ref(), b"30"); + assert_eq!(all[1].0.as_ref(), b"name"); + assert_eq!(all[1].1.as_ref(), b"alice"); +} + +#[tokio::test] +async fn hget_missing_field() { + let (_server, mut client) = connect().await; + client.hset("h", &[("exists", "yes")]).await.unwrap(); + assert!(client.hget("h", "missing").await.unwrap().is_none()); +} + +// --- sorted set commands --- + +#[tokio::test] +async fn zadd_zrange_zscore() { + let (_server, mut client) = connect().await; + client + .zadd("scores", &[(1.0, "alice"), (2.0, "bob"), (3.0, "carol")]) + .await + .unwrap(); + + let members = client.zrange("scores", 0, -1).await.unwrap(); + assert_eq!(members.len(), 3); + assert_eq!(members[0].as_ref(), b"alice"); + + let score = client.zscore("scores", "bob").await.unwrap(); + assert_eq!(score, Some(2.0)); +} + +#[tokio::test] +async fn zscore_missing_member_returns_none() { + let (_server, mut client) = connect().await; + client.zadd("zs", &[(1.0, "only")]).await.unwrap(); + let score = client.zscore("zs", "nobody").await.unwrap(); + assert!(score.is_none()); +} + +#[tokio::test] +async fn zrank_present_and_absent() { + let (_server, mut client) = connect().await; + client + .zadd("ranking", &[(10.0, "a"), (20.0, "b")]) + .await + .unwrap(); + assert_eq!(client.zrank("ranking", "a").await.unwrap(), Some(0)); + assert_eq!(client.zrank("ranking", "b").await.unwrap(), Some(1)); + assert!(client.zrank("ranking", "nobody").await.unwrap().is_none()); +} + +// --- pipeline --- + +#[tokio::test] +async fn pipeline_multiple_gets() { + let (_server, mut client) = connect().await; + client.set("p1", "v1").await.unwrap(); + client.set("p2", "v2").await.unwrap(); + client.set("p3", "v3").await.unwrap(); + + let frames = client + .execute_pipeline(Pipeline::new().get("p1").get("p2").get("p3")) + .await + .unwrap(); + + assert_eq!(frames.len(), 3); +} + +#[tokio::test] +async fn pipeline_empty_returns_empty_vec() { + let (_server, mut client) = connect().await; + let frames = client.execute_pipeline(Pipeline::new()).await.unwrap(); + assert!(frames.is_empty()); +} + +#[tokio::test] +async fn pipeline_mixed_commands() { + let (_server, mut client) = connect().await; + + let frames = client + .execute_pipeline( + Pipeline::new() + .set("batch_key", "batch_value") + .get("batch_key") + .incr("batch_counter"), + ) + .await + .unwrap(); + + assert_eq!(frames.len(), 3); +} + +// --- error surfacing --- + +#[tokio::test] +async fn wrongtype_error_surfaces_as_server_error() { + let (_server, mut client) = connect().await; + + // Store a list then try to GET it — server returns WRONGTYPE + client.rpush("mylist", &["item"]).await.unwrap(); + + let err = client.get("mylist").await.unwrap_err(); + assert!( + matches!(err, ClientError::Server(_)), + "expected Server error, got {err:?}" + ); +} diff --git a/tests/integration/src/main.rs b/tests/integration/src/main.rs index b9250156..70ef52cd 100644 --- a/tests/integration/src/main.rs +++ b/tests/integration/src/main.rs @@ -3,6 +3,7 @@ mod helpers; mod auth; mod basic_operations; mod cli; +mod client_typed_api; mod cluster; mod data_types; mod persistence; From 7a3b3a2a6220d70c2f1efd43c65ab16983ebe3ee Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Wed, 25 Feb 2026 12:58:05 -0500 Subject: [PATCH 2/5] feat(client): add 23 new typed commands to ember-client MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit closes the gap between the typed rust client and the server's full command set. mechanical additions, all following existing patterns. new decoders: - string_value: simple/bulk → String (TYPE, INFO, ECHO, BGSAVE) - float_value: bulk float string → f64 (INCRBYFLOAT) - scan_page: Array([cursor, [keys]]) → ScanPage (SCAN) - slowlog_entries: Array of entry arrays → Vec (SLOWLOG GET) - numsub_pairs: interleaved ch/count array → Vec<(Bytes, i64)> (PUBSUB NUMSUB) new public types: ScanPage, SlowlogEntry (re-exported from lib.rs) new Client methods: - strlen, incr_by_float - key_type, keys, rename, scan, pexpire - hmget - echo, unlink, info, bgsave, bgrewriteaof - slowlog_get, slowlog_len, slowlog_reset - publish, pubsub_channels, pubsub_numsub, pubsub_numpat new Pipeline builder methods for all new commands integration tests covering all new commands --- crates/ember-client/src/commands.rs | 522 ++++++++++++++++++++++ crates/ember-client/src/lib.rs | 1 + crates/ember-client/src/pipeline.rs | 60 +++ tests/integration/src/client_typed_api.rs | 215 ++++++++- 4 files changed, 797 insertions(+), 1 deletion(-) diff --git a/crates/ember-client/src/commands.rs b/crates/ember-client/src/commands.rs index 5761f1e2..cd1553d2 100644 --- a/crates/ember-client/src/commands.rs +++ b/crates/ember-client/src/commands.rs @@ -17,6 +17,32 @@ use ember_protocol::types::Frame; use crate::connection::{Client, ClientError}; use crate::pipeline::Pipeline; +// --- public types --- + +/// A page of keys returned by [`Client::scan`]. +/// +/// Iterate until `cursor` is `0` to walk the full keyspace. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ScanPage { + /// Cursor for the next call. `0` means iteration is complete. + pub cursor: u64, + /// Keys returned in this page. + pub keys: Vec, +} + +/// A single entry from [`Client::slowlog_get`]. +#[derive(Debug, Clone)] +pub struct SlowlogEntry { + /// Monotonically increasing log entry ID. + pub id: i64, + /// Unix timestamp (seconds) when the command was logged. + pub timestamp: i64, + /// Execution time in microseconds. + pub duration_us: i64, + /// The command and its arguments as raw bytes. + pub command: Vec, +} + // --- frame construction helpers --- fn cmd2(name: &'static [u8], a: &[u8]) -> Frame { @@ -282,6 +308,213 @@ fn scored_members(frame: Frame) -> Result, ClientError> { Ok(result) } +/// Decodes a simple string or bulk string response into a `String`. +/// +/// Used by: TYPE, INFO, ECHO, BGSAVE, BGREWRITEAOF. +fn string_value(frame: Frame) -> Result { + match frame { + Frame::Simple(s) => Ok(s), + Frame::Bulk(b) => String::from_utf8(b.to_vec()) + .map_err(|_| ClientError::Protocol("response is not valid UTF-8".into())), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected simple or bulk string, got {other:?}" + ))), + } +} + +/// Decodes a float returned as a bulk string. +/// +/// Used by: INCRBYFLOAT. +fn float_value(frame: Frame) -> Result { + match frame { + Frame::Bulk(b) => { + let s = std::str::from_utf8(&b) + .map_err(|_| ClientError::Protocol("float response is not valid UTF-8".into()))?; + s.parse::() + .map_err(|_| ClientError::Protocol(format!("not a valid float: {s:?}"))) + } + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected bulk float, got {other:?}" + ))), + } +} + +/// Decodes a SCAN response into a `ScanPage`. +/// +/// RESP3 layout: `Array([Bulk(cursor), Array([Bulk(key), ...])])`. +fn scan_page(frame: Frame) -> Result { + let elems = match frame { + Frame::Array(e) => e, + Frame::Error(e) => return Err(ClientError::Server(e)), + other => { + return Err(ClientError::Protocol(format!( + "expected array for SCAN, got {other:?}" + ))) + } + }; + + if elems.len() != 2 { + return Err(ClientError::Protocol(format!( + "SCAN response must have 2 elements, got {}", + elems.len() + ))); + } + + let mut iter = elems.into_iter(); + let cursor_frame = iter.next().unwrap(); + let keys_frame = iter.next().unwrap(); + + let cursor = match cursor_frame { + Frame::Bulk(b) => { + let s = std::str::from_utf8(&b) + .map_err(|_| ClientError::Protocol("SCAN cursor is not valid UTF-8".into()))?; + s.parse::() + .map_err(|_| ClientError::Protocol(format!("SCAN cursor is not a u64: {s:?}")))? + } + other => { + return Err(ClientError::Protocol(format!( + "expected bulk cursor in SCAN, got {other:?}" + ))) + } + }; + + let keys = bytes_vec(keys_frame)?; + Ok(ScanPage { cursor, keys }) +} + +/// Decodes a SLOWLOG GET response into a list of entries. +/// +/// Each entry: `Array([Integer(id), Integer(ts), Integer(us), Array([Bulk(cmd), ...])])`. +fn slowlog_entries(frame: Frame) -> Result, ClientError> { + let outer = match frame { + Frame::Array(e) => e, + Frame::Null => return Ok(Vec::new()), + Frame::Error(e) => return Err(ClientError::Server(e)), + other => { + return Err(ClientError::Protocol(format!( + "expected array for SLOWLOG, got {other:?}" + ))) + } + }; + + outer + .into_iter() + .map(|entry_frame| { + let entry = match entry_frame { + Frame::Array(e) => e, + other => { + return Err(ClientError::Protocol(format!( + "expected array for slowlog entry, got {other:?}" + ))) + } + }; + + if entry.len() < 4 { + return Err(ClientError::Protocol(format!( + "slowlog entry too short: {} elements", + entry.len() + ))); + } + + let id = match &entry[0] { + Frame::Integer(n) => *n, + other => { + return Err(ClientError::Protocol(format!( + "expected integer id in slowlog, got {other:?}" + ))) + } + }; + let timestamp = match &entry[1] { + Frame::Integer(n) => *n, + other => { + return Err(ClientError::Protocol(format!( + "expected integer timestamp in slowlog, got {other:?}" + ))) + } + }; + let duration_us = match &entry[2] { + Frame::Integer(n) => *n, + other => { + return Err(ClientError::Protocol(format!( + "expected integer duration in slowlog, got {other:?}" + ))) + } + }; + let command = match entry.into_iter().nth(3).unwrap() { + Frame::Array(parts) => parts + .into_iter() + .map(|p| match p { + Frame::Bulk(b) => Ok(b), + other => Err(ClientError::Protocol(format!( + "expected bulk in slowlog command, got {other:?}" + ))), + }) + .collect::, _>>()?, + other => { + return Err(ClientError::Protocol(format!( + "expected array for slowlog command, got {other:?}" + ))) + } + }; + + Ok(SlowlogEntry { + id, + timestamp, + duration_us, + command, + }) + }) + .collect() +} + +/// Decodes a PUBSUB NUMSUB response into channel/count pairs. +/// +/// Layout: `Array([Bulk(channel), Integer(count), ...])` — alternating. +fn numsub_pairs(frame: Frame) -> Result, ClientError> { + let elems = match frame { + Frame::Array(e) => e, + Frame::Null => return Ok(Vec::new()), + Frame::Error(e) => return Err(ClientError::Server(e)), + other => { + return Err(ClientError::Protocol(format!( + "expected array for PUBSUB NUMSUB, got {other:?}" + ))) + } + }; + + if elems.len() % 2 != 0 { + return Err(ClientError::Protocol(format!( + "PUBSUB NUMSUB array has odd length ({})", + elems.len() + ))); + } + + let mut result = Vec::with_capacity(elems.len() / 2); + let mut iter = elems.into_iter(); + while let (Some(ch_frame), Some(cnt_frame)) = (iter.next(), iter.next()) { + let channel = match ch_frame { + Frame::Bulk(b) => b, + other => { + return Err(ClientError::Protocol(format!( + "expected bulk channel in PUBSUB NUMSUB, got {other:?}" + ))) + } + }; + let count = match cnt_frame { + Frame::Integer(n) => n, + other => { + return Err(ClientError::Protocol(format!( + "expected integer count in PUBSUB NUMSUB, got {other:?}" + ))) + } + }; + result.push((channel, count)); + } + Ok(result) +} + // --- typed command API --- impl Client { @@ -728,6 +961,285 @@ impl Client { ok(frame) } + // --- more string commands --- + + /// Returns the length of the string at `key`. Returns 0 for missing keys. + pub async fn strlen(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"STRLEN", key.as_bytes())).await?; + integer(frame) + } + + /// Increments the float stored at `key` by `delta`. Returns the new value. + pub async fn incr_by_float(&mut self, key: &str, delta: f64) -> Result { + let d = delta.to_string(); + let frame = self + .send_frame(cmd3(b"INCRBYFLOAT", key.as_bytes(), d.as_bytes())) + .await?; + float_value(frame) + } + + // --- key commands --- + + /// Returns the type of the value stored at `key` as a string: + /// `"string"`, `"list"`, `"set"`, `"zset"`, `"hash"`, or `"none"`. + pub async fn key_type(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"TYPE", key.as_bytes())).await?; + string_value(frame) + } + + /// Returns all keys matching `pattern`. + /// + /// Use `"*"` to return every key. This is a blocking O(N) scan — prefer + /// [`Client::scan`] in production. + pub async fn keys(&mut self, pattern: &str) -> Result, ClientError> { + let frame = self.send_frame(cmd2(b"KEYS", pattern.as_bytes())).await?; + bytes_vec(frame) + } + + /// Renames `key` to `newkey`. Returns an error if `key` does not exist. + pub async fn rename(&mut self, key: &str, newkey: &str) -> Result<(), ClientError> { + let frame = self + .send_frame(cmd3(b"RENAME", key.as_bytes(), newkey.as_bytes())) + .await?; + ok(frame) + } + + /// Incrementally iterates keys in the keyspace. + /// + /// Pass `cursor: 0` to start a new iteration. Continue calling with the + /// returned cursor until the cursor is `0` again. An optional `pattern` + /// filters by glob and `count` hints at the page size (server may return + /// more or fewer). + pub async fn scan( + &mut self, + cursor: u64, + count: Option, + pattern: Option<&str>, + ) -> Result { + let cur = cursor.to_string(); + let mut parts = Vec::with_capacity(6); + parts.push(Frame::Bulk(Bytes::from_static(b"SCAN"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(cur.as_bytes()))); + if let Some(pat) = pattern { + parts.push(Frame::Bulk(Bytes::from_static(b"MATCH"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(pat.as_bytes()))); + } + if let Some(n) = count { + let ns = n.to_string(); + parts.push(Frame::Bulk(Bytes::from_static(b"COUNT"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(ns.as_bytes()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + scan_page(frame) + } + + /// Sets a timeout of `millis` milliseconds on `key`. Returns `true` if + /// the timeout was set, `false` if the key does not exist. + pub async fn pexpire(&mut self, key: &str, millis: u64) -> Result { + let ms = millis.to_string(); + let frame = self + .send_frame(cmd3(b"PEXPIRE", key.as_bytes(), ms.as_bytes())) + .await?; + bool_flag(frame) + } + + // --- more hash commands --- + + /// Returns values for multiple `fields` in the hash at `key`. Missing + /// fields are `None`. + pub async fn hmget( + &mut self, + key: &str, + fields: &[&str], + ) -> Result>, ClientError> { + let frame = self + .send_frame(cmd_key_and_keys(b"HMGET", key, fields)) + .await?; + optional_bytes_vec(frame) + } + + // --- more server commands --- + + /// Echoes `message` back from the server. Useful for round-trip testing. + pub async fn echo(&mut self, message: &str) -> Result { + let frame = self.send_frame(cmd2(b"ECHO", message.as_bytes())).await?; + match frame { + Frame::Bulk(b) => Ok(b), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected bulk for ECHO, got {other:?}" + ))), + } + } + + /// Deletes keys asynchronously. Behaves like `del` but frees memory in + /// the background for large values. Returns the number of keys removed. + pub async fn unlink(&mut self, keys: &[&str]) -> Result { + let frame = self + .send_frame(cmd_keys_only(b"UNLINK", keys)) + .await?; + integer(frame) + } + + /// Returns server information. Pass `Some("keyspace")` for a specific + /// section, or `None` for all sections. + pub async fn info(&mut self, section: Option<&str>) -> Result { + let frame = match section { + Some(s) => self.send_frame(cmd2(b"INFO", s.as_bytes())).await?, + None => { + self.send_frame(Frame::Array(vec![Frame::Bulk(Bytes::from_static(b"INFO"))])) + .await? + } + }; + match frame { + Frame::Bulk(b) => String::from_utf8(b.to_vec()) + .map_err(|_| ClientError::Protocol("INFO response is not valid UTF-8".into())), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected bulk for INFO, got {other:?}" + ))), + } + } + + /// Triggers a background snapshot (`BGSAVE`). Returns the server status + /// message. + pub async fn bgsave(&mut self) -> Result { + let frame = self + .send_frame(Frame::Array(vec![Frame::Bulk(Bytes::from_static(b"BGSAVE"))])) + .await?; + string_value(frame) + } + + /// Triggers an AOF rewrite in the background (`BGREWRITEAOF`). Returns + /// the server status message. + pub async fn bgrewriteaof(&mut self) -> Result { + let frame = self + .send_frame(Frame::Array(vec![Frame::Bulk(Bytes::from_static( + b"BGREWRITEAOF", + ))])) + .await?; + string_value(frame) + } + + // --- slowlog commands --- + + /// Returns up to `count` recent slow-log entries, or all entries if + /// `count` is `None`. + pub async fn slowlog_get( + &mut self, + count: Option, + ) -> Result, ClientError> { + let frame = match count { + Some(n) => { + let ns = n.to_string(); + self.send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"SLOWLOG")), + Frame::Bulk(Bytes::from_static(b"GET")), + Frame::Bulk(Bytes::copy_from_slice(ns.as_bytes())), + ])) + .await? + } + None => { + self.send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"SLOWLOG")), + Frame::Bulk(Bytes::from_static(b"GET")), + ])) + .await? + } + }; + slowlog_entries(frame) + } + + /// Returns the number of entries currently in the slow log. + pub async fn slowlog_len(&mut self) -> Result { + let frame = self + .send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"SLOWLOG")), + Frame::Bulk(Bytes::from_static(b"LEN")), + ])) + .await?; + integer(frame) + } + + /// Clears all entries from the slow log. + pub async fn slowlog_reset(&mut self) -> Result<(), ClientError> { + let frame = self + .send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"SLOWLOG")), + Frame::Bulk(Bytes::from_static(b"RESET")), + ])) + .await?; + ok(frame) + } + + // --- pub/sub commands (request-response only) --- + + /// Publishes `message` to `channel`. Returns the number of subscribers + /// that received the message. + pub async fn publish( + &mut self, + channel: &str, + message: impl AsRef<[u8]>, + ) -> Result { + let frame = self + .send_frame(cmd3(b"PUBLISH", channel.as_bytes(), message.as_ref())) + .await?; + integer(frame) + } + + /// Returns the names of active pub/sub channels. Pass `Some(pattern)` to + /// filter by glob, or `None` for all channels. + pub async fn pubsub_channels( + &mut self, + pattern: Option<&str>, + ) -> Result, ClientError> { + let frame = match pattern { + Some(p) => { + self.send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"PUBSUB")), + Frame::Bulk(Bytes::from_static(b"CHANNELS")), + Frame::Bulk(Bytes::copy_from_slice(p.as_bytes())), + ])) + .await? + } + None => { + self.send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"PUBSUB")), + Frame::Bulk(Bytes::from_static(b"CHANNELS")), + ])) + .await? + } + }; + bytes_vec(frame) + } + + /// Returns the subscriber counts for the given channels as + /// `(channel, count)` pairs. + pub async fn pubsub_numsub( + &mut self, + channels: &[&str], + ) -> Result, ClientError> { + let mut parts = Vec::with_capacity(2 + channels.len()); + parts.push(Frame::Bulk(Bytes::from_static(b"PUBSUB"))); + parts.push(Frame::Bulk(Bytes::from_static(b"NUMSUB"))); + for ch in channels { + parts.push(Frame::Bulk(Bytes::copy_from_slice(ch.as_bytes()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + numsub_pairs(frame) + } + + /// Returns the number of active pattern subscriptions across all clients. + pub async fn pubsub_numpat(&mut self) -> Result { + let frame = self + .send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"PUBSUB")), + Frame::Bulk(Bytes::from_static(b"NUMPAT")), + ])) + .await?; + integer(frame) + } + // --- pipeline --- /// Executes all commands queued in `pipeline` as a single batch. @@ -772,6 +1284,16 @@ fn cmd_key_values>(cmd: &'static [u8], key: &str, values: &[V]) - Frame::Array(parts) } +/// `CMD key1 key2 ...` (no separate leading key argument) +fn cmd_keys_only(cmd: &'static [u8], keys: &[&str]) -> Frame { + let mut parts = Vec::with_capacity(1 + keys.len()); + parts.push(Frame::Bulk(Bytes::from_static(cmd))); + for k in keys { + parts.push(Frame::Bulk(Bytes::copy_from_slice(k.as_bytes()))); + } + Frame::Array(parts) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/ember-client/src/lib.rs b/crates/ember-client/src/lib.rs index a729d466..a107ca76 100644 --- a/crates/ember-client/src/lib.rs +++ b/crates/ember-client/src/lib.rs @@ -51,6 +51,7 @@ mod pipeline; #[cfg(feature = "tls")] pub mod tls; +pub use commands::{ScanPage, SlowlogEntry}; pub use connection::{Client, ClientError}; pub use ember_protocol::types::Frame; pub use pipeline::Pipeline; diff --git a/crates/ember-client/src/pipeline.rs b/crates/ember-client/src/pipeline.rs index 0346d40d..4a3f4039 100644 --- a/crates/ember-client/src/pipeline.rs +++ b/crates/ember-client/src/pipeline.rs @@ -210,6 +210,66 @@ impl Pipeline { self.push(array3(b"ZSCORE", key.as_bytes(), member.as_bytes())) } + // --- more string commands --- + + /// Queues a `STRLEN key` command. + pub fn strlen(self, key: &str) -> Self { + self.push(array2(b"STRLEN", key.as_bytes())) + } + + /// Queues an `INCRBYFLOAT key delta` command. + pub fn incr_by_float(self, key: &str, delta: f64) -> Self { + let d = delta.to_string(); + self.push(array3(b"INCRBYFLOAT", key.as_bytes(), d.as_bytes())) + } + + // --- key commands --- + + /// Queues a `TYPE key` command. + pub fn key_type(self, key: &str) -> Self { + self.push(array2(b"TYPE", key.as_bytes())) + } + + /// Queues a `KEYS pattern` command. + pub fn keys(self, pattern: &str) -> Self { + self.push(array2(b"KEYS", pattern.as_bytes())) + } + + /// Queues a `RENAME key newkey` command. + pub fn rename(self, key: &str, newkey: &str) -> Self { + self.push(array3(b"RENAME", key.as_bytes(), newkey.as_bytes())) + } + + /// Queues a `PEXPIRE key millis` command. + pub fn pexpire(self, key: &str, millis: u64) -> Self { + let ms = millis.to_string(); + self.push(array3(b"PEXPIRE", key.as_bytes(), ms.as_bytes())) + } + + /// Queues an `UNLINK key [key ...]` command. + pub fn unlink(self, keys: &[&str]) -> Self { + self.push(array_with_keys(b"UNLINK", keys)) + } + + // --- more hash commands --- + + /// Queues an `HMGET key field [field ...]` command. + pub fn hmget(self, key: &str, fields: &[&str]) -> Self { + self.push(array_with_key_and_keys(b"HMGET", key, fields)) + } + + // --- server commands --- + + /// Queues an `ECHO message` command. + pub fn echo(self, message: &str) -> Self { + self.push(array2(b"ECHO", message.as_bytes())) + } + + /// Queues a `PUBLISH channel message` command. + pub fn publish(self, channel: &str, message: impl AsRef<[u8]>) -> Self { + self.push(array3(b"PUBLISH", channel.as_bytes(), message.as_ref())) + } + // --- internal --- fn push(mut self, frame: Frame) -> Self { diff --git a/tests/integration/src/client_typed_api.rs b/tests/integration/src/client_typed_api.rs index 170ac479..7a73f066 100644 --- a/tests/integration/src/client_typed_api.rs +++ b/tests/integration/src/client_typed_api.rs @@ -4,7 +4,7 @@ //! connects with the typed `Client`. These tests exercise the full stack — //! TCP, RESP3 framing, command dispatch, and response decoding. -use ember_client::{Client, ClientError, Pipeline}; +use ember_client::{Client, ClientError, Pipeline, ScanPage}; use crate::helpers::TestServer; @@ -226,6 +226,219 @@ async fn pipeline_mixed_commands() { assert_eq!(frames.len(), 3); } +// --- strlen --- + +#[tokio::test] +async fn strlen_existing_and_missing() { + let (_server, mut client) = connect().await; + client.set("slen", "hello").await.unwrap(); + assert_eq!(client.strlen("slen").await.unwrap(), 5); + assert_eq!(client.strlen("nosuchkey").await.unwrap(), 0); +} + +// --- incr_by_float --- + +#[tokio::test] +async fn incr_by_float_positive_and_negative() { + let (_server, mut client) = connect().await; + client.set("flt", "10.5").await.unwrap(); + let v = client.incr_by_float("flt", 1.5).await.unwrap(); + assert!((v - 12.0).abs() < 1e-9); + let v2 = client.incr_by_float("flt", -2.0).await.unwrap(); + assert!((v2 - 10.0).abs() < 1e-9); +} + +// --- pexpire --- + +#[tokio::test] +async fn pexpire_sets_ttl_in_millis() { + let (_server, mut client) = connect().await; + client.set("px_key", "v").await.unwrap(); + let set = client.pexpire("px_key", 60_000).await.unwrap(); + assert!(set); + let pttl = client.pttl("px_key").await.unwrap(); + assert!(pttl > 0 && pttl <= 60_000); +} + +// --- hmget --- + +#[tokio::test] +async fn hmget_subset_and_missing() { + let (_server, mut client) = connect().await; + client + .hset("hm", &[("a", "1"), ("b", "2"), ("c", "3")]) + .await + .unwrap(); + let vals = client.hmget("hm", &["a", "c", "missing"]).await.unwrap(); + assert_eq!(vals.len(), 3); + assert_eq!(vals[0].as_deref(), Some(b"1" as &[u8])); + assert_eq!(vals[1].as_deref(), Some(b"3" as &[u8])); + assert!(vals[2].is_none()); +} + +// --- key_type --- + +#[tokio::test] +async fn key_type_returns_correct_types() { + let (_server, mut client) = connect().await; + client.set("str_key", "v").await.unwrap(); + assert_eq!(client.key_type("str_key").await.unwrap(), "string"); + + client.rpush("lst_key", &["item"]).await.unwrap(); + assert_eq!(client.key_type("lst_key").await.unwrap(), "list"); + + assert_eq!(client.key_type("no_such_key").await.unwrap(), "none"); +} + +// --- keys --- + +#[tokio::test] +async fn keys_glob_returns_matching_keys() { + let (_server, mut client) = connect().await; + client.set("pfx:a", "1").await.unwrap(); + client.set("pfx:b", "2").await.unwrap(); + client.set("other", "3").await.unwrap(); + let mut matched = client.keys("pfx:*").await.unwrap(); + matched.sort(); + assert_eq!(matched.len(), 2); + assert_eq!(matched[0].as_ref(), b"pfx:a"); + assert_eq!(matched[1].as_ref(), b"pfx:b"); +} + +// --- rename --- + +#[tokio::test] +async fn rename_then_old_key_gone() { + let (_server, mut client) = connect().await; + client.set("old_name", "value").await.unwrap(); + client.rename("old_name", "new_name").await.unwrap(); + assert!(client.get("old_name").await.unwrap().is_none()); + assert_eq!( + client.get("new_name").await.unwrap().as_deref(), + Some(b"value" as &[u8]) + ); +} + +// --- scan --- + +#[tokio::test] +async fn scan_full_iteration_collects_all_keys() { + let (_server, mut client) = connect().await; + let expected_keys = ["scan_a", "scan_b", "scan_c"]; + for k in &expected_keys { + client.set(k, "v").await.unwrap(); + } + + let mut all_keys = Vec::new(); + let mut cursor = 0u64; + loop { + let page = client.scan(cursor, None, Some("scan_*")).await.unwrap(); + all_keys.extend(page.keys); + cursor = page.cursor; + if cursor == 0 { + break; + } + } + + all_keys.sort(); + assert_eq!(all_keys.len(), 3); + assert_eq!(all_keys[0].as_ref(), b"scan_a"); + assert_eq!(all_keys[2].as_ref(), b"scan_c"); +} + +// --- echo --- + +#[tokio::test] +async fn echo_returns_same_bytes() { + let (_server, mut client) = connect().await; + let reply = client.echo("hello world").await.unwrap(); + assert_eq!(reply.as_ref(), b"hello world"); +} + +// --- unlink --- + +#[tokio::test] +async fn unlink_returns_count_like_del() { + let (_server, mut client) = connect().await; + client.set("ul1", "v").await.unwrap(); + client.set("ul2", "v").await.unwrap(); + let removed = client.unlink(&["ul1", "ul2", "ul_missing"]).await.unwrap(); + assert_eq!(removed, 2); + assert!(client.get("ul1").await.unwrap().is_none()); +} + +// --- info --- + +#[tokio::test] +async fn info_returns_non_empty_string() { + let (_server, mut client) = connect().await; + let output = client.info(None).await.unwrap(); + assert!(!output.is_empty()); +} + +// --- bgsave --- + +#[tokio::test] +async fn bgsave_returns_status_string() { + let (_server, mut client) = connect().await; + let status = client.bgsave().await.unwrap(); + assert!(!status.is_empty()); +} + +// --- slowlog --- + +#[tokio::test] +async fn slowlog_len_is_non_negative() { + let (_server, mut client) = connect().await; + let len = client.slowlog_len().await.unwrap(); + assert!(len >= 0); +} + +#[tokio::test] +async fn slowlog_reset_clears_log() { + let (_server, mut client) = connect().await; + client.slowlog_reset().await.unwrap(); + assert_eq!(client.slowlog_len().await.unwrap(), 0); +} + +// --- pub/sub (request-response) --- + +#[tokio::test] +async fn publish_to_channel_with_no_subscribers_returns_zero() { + let (_server, mut client) = connect().await; + let count = client.publish("unsubscribed_channel", "msg").await.unwrap(); + assert_eq!(count, 0); +} + +#[tokio::test] +async fn pubsub_channels_with_pattern_and_none() { + let (_server, mut client) = connect().await; + // no subscribers active — both should return empty + let with_pat = client.pubsub_channels(Some("*")).await.unwrap(); + let all = client.pubsub_channels(None).await.unwrap(); + assert!(with_pat.is_empty()); + assert!(all.is_empty()); +} + +#[tokio::test] +async fn pubsub_numsub_empty_channels() { + let (_server, mut client) = connect().await; + let pairs = client + .pubsub_numsub(&["no_one_here"]) + .await + .unwrap(); + assert_eq!(pairs.len(), 1); + assert_eq!(pairs[0].0.as_ref(), b"no_one_here"); + assert_eq!(pairs[0].1, 0); +} + +#[tokio::test] +async fn pubsub_numpat_is_zero_before_subscriptions() { + let (_server, mut client) = connect().await; + let n = client.pubsub_numpat().await.unwrap(); + assert_eq!(n, 0); +} + // --- error surfacing --- #[tokio::test] From e3bcbad3dcb132b2d07c37c117b917e78e022d93 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Wed, 25 Feb 2026 13:00:13 -0500 Subject: [PATCH 3/5] feat(client): add Subscriber type for pub/sub mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit adds a dedicated subscriber connection type that wraps a Client once it enters pub/sub mode. regular request-response commands cannot be issued on the same connection while subscribed — callers use a separate Client. new types: Subscriber, Message (re-exported from lib.rs) new Client methods: - subscribe(channels) → Subscriber - psubscribe(patterns) → Subscriber Subscriber methods: - recv() — waits for the next message, skips confirmation frames - subscribe(channels) — add more channels without leaving sub mode - unsubscribe(channels) → Client — remove channels and reclaim connection connection.rs: added write_frame() and pub(crate) read_response() so Subscriber can drive the connection directly integration tests: subscriber_recv_message, subscriber_multiple_channels, psubscriber_pattern_match --- crates/ember-client/src/commands.rs | 37 +++++ crates/ember-client/src/connection.rs | 14 +- crates/ember-client/src/lib.rs | 2 + crates/ember-client/src/subscriber.rs | 185 ++++++++++++++++++++++ tests/integration/src/client_typed_api.rs | 60 ++++++- 5 files changed, 296 insertions(+), 2 deletions(-) create mode 100644 crates/ember-client/src/subscriber.rs diff --git a/crates/ember-client/src/commands.rs b/crates/ember-client/src/commands.rs index cd1553d2..2ab3a10f 100644 --- a/crates/ember-client/src/commands.rs +++ b/crates/ember-client/src/commands.rs @@ -16,6 +16,7 @@ use ember_protocol::types::Frame; use crate::connection::{Client, ClientError}; use crate::pipeline::Pipeline; +use crate::subscriber::Subscriber; // --- public types --- @@ -1240,6 +1241,42 @@ impl Client { integer(frame) } + // --- pub/sub subscriber mode --- + + /// Puts the connection into subscriber mode on `channels`. + /// + /// The connection is consumed and a [`Subscriber`] is returned. Use a + /// separate [`Client`] for regular commands while subscribed. + pub async fn subscribe(mut self, channels: &[&str]) -> Result { + let mut parts = Vec::with_capacity(1 + channels.len()); + parts.push(Frame::Bulk(Bytes::from_static(b"SUBSCRIBE"))); + for ch in channels { + parts.push(Frame::Bulk(Bytes::copy_from_slice(ch.as_bytes()))); + } + self.write_frame(Frame::Array(parts)).await?; + // drain confirmation frames (one per channel) + for _ in 0..channels.len() { + self.read_response().await?; + } + Ok(Subscriber::new(self)) + } + + /// Same as [`subscribe`](Client::subscribe) but subscribes to glob + /// patterns with `PSUBSCRIBE`. + pub async fn psubscribe(mut self, patterns: &[&str]) -> Result { + let mut parts = Vec::with_capacity(1 + patterns.len()); + parts.push(Frame::Bulk(Bytes::from_static(b"PSUBSCRIBE"))); + for p in patterns { + parts.push(Frame::Bulk(Bytes::copy_from_slice(p.as_bytes()))); + } + self.write_frame(Frame::Array(parts)).await?; + // drain confirmation frames (one per pattern) + for _ in 0..patterns.len() { + self.read_response().await?; + } + Ok(Subscriber::new(self)) + } + // --- pipeline --- /// Executes all commands queued in `pipeline` as a single batch. diff --git a/crates/ember-client/src/connection.rs b/crates/ember-client/src/connection.rs index d01a3c67..09fdb551 100644 --- a/crates/ember-client/src/connection.rs +++ b/crates/ember-client/src/connection.rs @@ -242,8 +242,20 @@ impl Client { let _ = self.transport.shutdown().await; } + /// Serializes and writes a single frame without waiting for a response. + /// + /// Used by [`Subscriber`] to send SUBSCRIBE/UNSUBSCRIBE frames before + /// draining the confirmation responses separately. + pub(crate) async fn write_frame(&mut self, frame: Frame) -> Result<(), ClientError> { + self.write_buf.clear(); + frame.serialize(&mut self.write_buf); + self.transport.write_all(&self.write_buf).await?; + self.transport.flush().await?; + Ok(()) + } + /// Reads a complete RESP3 frame from the server. - async fn read_response(&mut self) -> Result { + pub(crate) async fn read_response(&mut self) -> Result { loop { if !self.read_buf.is_empty() { match parse_frame(&self.read_buf) { diff --git a/crates/ember-client/src/lib.rs b/crates/ember-client/src/lib.rs index a107ca76..d4acd7bc 100644 --- a/crates/ember-client/src/lib.rs +++ b/crates/ember-client/src/lib.rs @@ -48,6 +48,7 @@ mod commands; mod connection; mod pipeline; +pub mod subscriber; #[cfg(feature = "tls")] pub mod tls; @@ -55,5 +56,6 @@ pub use commands::{ScanPage, SlowlogEntry}; pub use connection::{Client, ClientError}; pub use ember_protocol::types::Frame; pub use pipeline::Pipeline; +pub use subscriber::{Message, Subscriber}; #[cfg(feature = "tls")] pub use tls::TlsClientConfig; diff --git a/crates/ember-client/src/subscriber.rs b/crates/ember-client/src/subscriber.rs new file mode 100644 index 00000000..7f7a714a --- /dev/null +++ b/crates/ember-client/src/subscriber.rs @@ -0,0 +1,185 @@ +//! Pub/sub subscriber mode. +//! +//! When a connection issues [`Client::subscribe`] or [`Client::psubscribe`] it +//! enters pub/sub mode: the server will push [`Message`] frames whenever a +//! matching publish event occurs. Normal request-response commands cannot be +//! issued on the same connection while it is in sub mode. +//! +//! Use a separate [`Client`] for regular commands while subscribed. + +use bytes::Bytes; +use ember_protocol::types::Frame; + +use crate::connection::{Client, ClientError}; + +/// A message pushed by the server to a subscribed connection. +#[derive(Debug, Clone)] +pub struct Message { + /// The channel the message was published to. + pub channel: Bytes, + /// The message payload. + pub data: Bytes, + /// Set only for pattern-matched messages (`PSUBSCRIBE`). Contains the + /// pattern that matched the channel. + pub pattern: Option, +} + +/// A connection locked into pub/sub mode. +/// +/// Obtained by calling [`Client::subscribe`] or [`Client::psubscribe`]. +/// The underlying transport is consumed — create a separate [`Client`] +/// for regular commands while this subscriber is active. +/// +/// # Example +/// +/// ```no_run +/// use ember_client::Client; +/// +/// #[tokio::main] +/// async fn main() -> Result<(), ember_client::ClientError> { +/// let mut publisher = Client::connect("127.0.0.1", 6379).await?; +/// let subscriber_conn = Client::connect("127.0.0.1", 6379).await?; +/// +/// let mut sub = subscriber_conn.subscribe(&["news"]).await?; +/// +/// publisher.publish("news", "breaking: hello world").await?; +/// +/// let msg = sub.recv().await?; +/// println!("got: {:?}", msg.data); +/// Ok(()) +/// } +/// ``` +pub struct Subscriber { + inner: Client, +} + +impl Subscriber { + pub(crate) fn new(inner: Client) -> Self { + Self { inner } + } + + /// Blocks until the next message arrives on any subscribed channel. + /// + /// Subscription confirmation frames (`subscribe`/`psubscribe`) are + /// skipped silently — only actual message frames are returned. + pub async fn recv(&mut self) -> Result { + loop { + let frame = self.inner.read_response().await?; + if let Some(msg) = try_parse_message(frame)? { + return Ok(msg); + } + // confirmation frame (subscribe/unsubscribe/psubscribe/punsubscribe) + // — loop back and wait for the next frame + } + } + + /// Subscribes to additional channels without leaving sub mode. + pub async fn subscribe(&mut self, channels: &[&str]) -> Result<(), ClientError> { + let mut parts = Vec::with_capacity(1 + channels.len()); + parts.push(Frame::Bulk(Bytes::from_static(b"SUBSCRIBE"))); + for ch in channels { + parts.push(Frame::Bulk(Bytes::copy_from_slice(ch.as_bytes()))); + } + self.inner.write_frame(Frame::Array(parts)).await?; + // drain the confirmation frames (one per channel) + for _ in 0..channels.len() { + self.inner.read_response().await?; + } + Ok(()) + } + + /// Unsubscribes from the given channels. + /// + /// When all subscriptions have been removed the inner [`Client`] is + /// returned so the connection can be reused for regular commands. + pub async fn unsubscribe(mut self, channels: &[&str]) -> Result { + let mut parts = Vec::with_capacity(1 + channels.len()); + parts.push(Frame::Bulk(Bytes::from_static(b"UNSUBSCRIBE"))); + for ch in channels { + parts.push(Frame::Bulk(Bytes::copy_from_slice(ch.as_bytes()))); + } + self.inner.write_frame(Frame::Array(parts)).await?; + // drain the unsubscribe confirmation frames + for _ in 0..channels.len() { + self.inner.read_response().await?; + } + Ok(self.inner) + } +} + +/// Tries to parse a push frame into a [`Message`]. +/// +/// Returns `Ok(Some(_))` for `message` and `pmessage` frames. +/// Returns `Ok(None)` for subscription management frames. +/// Returns `Err` for unexpected or malformed frames. +fn try_parse_message(frame: Frame) -> Result, ClientError> { + let elems = match frame { + Frame::Array(e) => e, + Frame::Error(e) => return Err(ClientError::Server(e)), + other => { + return Err(ClientError::Protocol(format!( + "expected array push frame, got {other:?}" + ))) + } + }; + + if elems.len() < 3 { + return Err(ClientError::Protocol(format!( + "push frame too short: {} elements", + elems.len() + ))); + } + + let kind = match &elems[0] { + Frame::Bulk(b) => b.clone(), + Frame::Simple(s) => Bytes::copy_from_slice(s.as_bytes()), + other => { + return Err(ClientError::Protocol(format!( + "expected bulk/simple frame kind, got {other:?}" + ))) + } + }; + + match kind.as_ref() { + b"message" => { + if elems.len() < 3 { + return Err(ClientError::Protocol( + "message frame has fewer than 3 elements".into(), + )); + } + let channel = bulk_bytes(elems[1].clone())?; + let data = bulk_bytes(elems[2].clone())?; + Ok(Some(Message { + channel, + data, + pattern: None, + })) + } + b"pmessage" => { + if elems.len() < 4 { + return Err(ClientError::Protocol( + "pmessage frame has fewer than 4 elements".into(), + )); + } + let pattern = bulk_bytes(elems[1].clone())?; + let channel = bulk_bytes(elems[2].clone())?; + let data = bulk_bytes(elems[3].clone())?; + Ok(Some(Message { + channel, + data, + pattern: Some(pattern), + })) + } + // subscribe / unsubscribe / psubscribe / punsubscribe confirmations + _ => Ok(None), + } +} + +fn bulk_bytes(frame: Frame) -> Result { + match frame { + Frame::Bulk(b) => Ok(b), + other => Err(ClientError::Protocol(format!( + "expected bulk in push frame, got {other:?}" + ))), + } +} diff --git a/tests/integration/src/client_typed_api.rs b/tests/integration/src/client_typed_api.rs index 7a73f066..5c42a44b 100644 --- a/tests/integration/src/client_typed_api.rs +++ b/tests/integration/src/client_typed_api.rs @@ -4,7 +4,7 @@ //! connects with the typed `Client`. These tests exercise the full stack — //! TCP, RESP3 framing, command dispatch, and response decoding. -use ember_client::{Client, ClientError, Pipeline, ScanPage}; +use ember_client::{Client, ClientError, Pipeline}; use crate::helpers::TestServer; @@ -439,6 +439,64 @@ async fn pubsub_numpat_is_zero_before_subscriptions() { assert_eq!(n, 0); } +// --- subscriber --- + +#[tokio::test] +async fn subscriber_recv_message() { + let server = crate::helpers::TestServer::start(); + let sub_conn = Client::connect("127.0.0.1", server.port) + .await + .expect("connect failed"); + let mut publisher = Client::connect("127.0.0.1", server.port) + .await + .expect("connect failed"); + + let mut sub = sub_conn.subscribe(&["typed_events"]).await.unwrap(); + publisher.publish("typed_events", "hello").await.unwrap(); + + let msg = sub.recv().await.unwrap(); + assert_eq!(msg.channel.as_ref(), b"typed_events"); + assert_eq!(msg.data.as_ref(), b"hello"); + assert!(msg.pattern.is_none()); +} + +#[tokio::test] +async fn subscriber_multiple_channels_correct_message() { + let server = crate::helpers::TestServer::start(); + let sub_conn = Client::connect("127.0.0.1", server.port) + .await + .expect("connect failed"); + let mut publisher = Client::connect("127.0.0.1", server.port) + .await + .expect("connect failed"); + + let mut sub = sub_conn.subscribe(&["ch_a", "ch_b"]).await.unwrap(); + publisher.publish("ch_b", "payload").await.unwrap(); + + let msg = sub.recv().await.unwrap(); + assert_eq!(msg.channel.as_ref(), b"ch_b"); + assert_eq!(msg.data.as_ref(), b"payload"); +} + +#[tokio::test] +async fn psubscriber_pattern_match() { + let server = crate::helpers::TestServer::start(); + let sub_conn = Client::connect("127.0.0.1", server.port) + .await + .expect("connect failed"); + let mut publisher = Client::connect("127.0.0.1", server.port) + .await + .expect("connect failed"); + + let mut sub = sub_conn.psubscribe(&["typed:*"]).await.unwrap(); + publisher.publish("typed:update", "data").await.unwrap(); + + let msg = sub.recv().await.unwrap(); + assert_eq!(msg.channel.as_ref(), b"typed:update"); + assert_eq!(msg.data.as_ref(), b"data"); + assert_eq!(msg.pattern.as_deref(), Some(b"typed:*" as &[u8])); +} + // --- error surfacing --- #[tokio::test] From 4abb06f710ce221c71cf15c4bef90236ddc0235b Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Wed, 25 Feb 2026 13:02:57 -0500 Subject: [PATCH 4/5] feat(client): vector set commands behind 'vector' feature flag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit adds VADD, VADD_BATCH, VSIM, VREM, VGET, VCARD, VDIM, VINFO to the typed client API, gated behind the new 'vector' cargo feature. new public type: SimResult { element: Bytes, distance: f32 } new Client methods (cfg(feature = "vector")): - vadd(key, element, vector) → bool - vadd_batch(key, dim, entries) → i64 - vsim(key, query, count) → Vec (always WITHSCORES) - vrem(key, element) → bool - vget(key, element) → Option> - vcard(key) → i64 - vdim(key) → i64 - vinfo(key) → Vec<(Bytes, Bytes)> floats are encoded as decimal strings in RESP3, matching the server's wire format. VADD_BATCH uses the text mode (DIM n then element floats); binary blob mode is not exposed since it's a platform-specific concern. --- crates/ember-client/Cargo.toml | 3 + crates/ember-client/src/lib.rs | 4 + crates/ember-client/src/vector.rs | 338 ++++++++++++++++++++++++++++++ 3 files changed, 345 insertions(+) create mode 100644 crates/ember-client/src/vector.rs diff --git a/crates/ember-client/Cargo.toml b/crates/ember-client/Cargo.toml index 17ae56ce..0a82db6f 100644 --- a/crates/ember-client/Cargo.toml +++ b/crates/ember-client/Cargo.toml @@ -30,3 +30,6 @@ tls = [ "dep:rustls-pki-types", "dep:rustls-native-certs", ] +# Vector set commands (VADD, VSIM, etc.). Requires the server to be built +# with the same feature enabled. +vector = [] diff --git a/crates/ember-client/src/lib.rs b/crates/ember-client/src/lib.rs index d4acd7bc..178f3c6c 100644 --- a/crates/ember-client/src/lib.rs +++ b/crates/ember-client/src/lib.rs @@ -51,6 +51,8 @@ mod pipeline; pub mod subscriber; #[cfg(feature = "tls")] pub mod tls; +#[cfg(feature = "vector")] +pub mod vector; pub use commands::{ScanPage, SlowlogEntry}; pub use connection::{Client, ClientError}; @@ -59,3 +61,5 @@ pub use pipeline::Pipeline; pub use subscriber::{Message, Subscriber}; #[cfg(feature = "tls")] pub use tls::TlsClientConfig; +#[cfg(feature = "vector")] +pub use vector::SimResult; diff --git a/crates/ember-client/src/vector.rs b/crates/ember-client/src/vector.rs new file mode 100644 index 00000000..2b4a3a26 --- /dev/null +++ b/crates/ember-client/src/vector.rs @@ -0,0 +1,338 @@ +//! Vector set commands (`VADD`, `VSIM`, etc.). +//! +//! Enabled by the `vector` cargo feature. Requires the server to be built +//! with the same feature enabled. + +use bytes::Bytes; +use ember_protocol::types::Frame; + +use crate::connection::{Client, ClientError}; + +// --- public types --- + +/// One result from a [`Client::vsim`] nearest-neighbour search. +#[derive(Debug, Clone)] +pub struct SimResult { + /// The element name. + pub element: Bytes, + /// Distance from the query vector. Lower is closer (metric-dependent). + pub distance: f32, +} + +// --- decoders --- + +fn vsim_results(frame: Frame, with_scores: bool) -> Result, ClientError> { + let elems = match frame { + Frame::Array(e) => e, + Frame::Null => return Ok(Vec::new()), + Frame::Error(e) => return Err(ClientError::Server(e)), + other => { + return Err(ClientError::Protocol(format!( + "expected array for VSIM, got {other:?}" + ))) + } + }; + + if !with_scores { + return elems + .into_iter() + .map(|e| match e { + Frame::Bulk(b) => Ok(SimResult { + element: b, + distance: 0.0, + }), + other => Err(ClientError::Protocol(format!( + "expected bulk element in VSIM, got {other:?}" + ))), + }) + .collect(); + } + + if elems.len() % 2 != 0 { + return Err(ClientError::Protocol(format!( + "VSIM WITHSCORES array has odd length ({})", + elems.len() + ))); + } + + let mut result = Vec::with_capacity(elems.len() / 2); + let mut iter = elems.into_iter(); + while let (Some(elem_frame), Some(dist_frame)) = (iter.next(), iter.next()) { + let element = match elem_frame { + Frame::Bulk(b) => b, + other => { + return Err(ClientError::Protocol(format!( + "expected bulk element in VSIM WITHSCORES, got {other:?}" + ))) + } + }; + let distance = match dist_frame { + Frame::Bulk(b) => { + let s = std::str::from_utf8(&b).map_err(|_| { + ClientError::Protocol("VSIM distance is not valid UTF-8".into()) + })?; + s.parse::().map_err(|_| { + ClientError::Protocol(format!("VSIM distance is not a valid float: {s:?}")) + })? + } + other => { + return Err(ClientError::Protocol(format!( + "expected bulk distance in VSIM WITHSCORES, got {other:?}" + ))) + } + }; + result.push(SimResult { element, distance }); + } + Ok(result) +} + +fn vget_vector(frame: Frame) -> Result>, ClientError> { + match frame { + Frame::Array(elems) => { + let floats = elems + .into_iter() + .map(|e| match e { + Frame::Bulk(b) => { + let s = std::str::from_utf8(&b).map_err(|_| { + ClientError::Protocol("vector component is not valid UTF-8".into()) + })?; + s.parse::().map_err(|_| { + ClientError::Protocol(format!( + "vector component is not a valid float: {s:?}" + )) + }) + } + other => Err(ClientError::Protocol(format!( + "expected bulk float in VGET response, got {other:?}" + ))), + }) + .collect::, _>>()?; + Ok(Some(floats)) + } + Frame::Null => Ok(None), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected array or null for VGET, got {other:?}" + ))), + } +} + +fn vinfo_pairs(frame: Frame) -> Result, ClientError> { + match frame { + Frame::Array(elems) => { + if elems.len() % 2 != 0 { + return Err(ClientError::Protocol(format!( + "VINFO array has odd length ({})", + elems.len() + ))); + } + let mut result = Vec::with_capacity(elems.len() / 2); + let mut iter = elems.into_iter(); + while let (Some(k), Some(v)) = (iter.next(), iter.next()) { + let key = match k { + Frame::Bulk(b) => b, + Frame::Simple(s) => Bytes::copy_from_slice(s.as_bytes()), + other => { + return Err(ClientError::Protocol(format!( + "expected bulk key in VINFO, got {other:?}" + ))) + } + }; + let val = match v { + Frame::Bulk(b) => b, + Frame::Integer(n) => Bytes::copy_from_slice(n.to_string().as_bytes()), + other => { + return Err(ClientError::Protocol(format!( + "expected bulk/integer value in VINFO, got {other:?}" + ))) + } + }; + result.push((key, val)); + } + Ok(result) + } + Frame::Null => Ok(Vec::new()), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected array for VINFO, got {other:?}" + ))), + } +} + +// --- Client methods --- + +impl Client { + /// Adds `element` with `vector` to the vector set at `key`. Returns `true` + /// if the element was newly inserted, `false` if it already existed and + /// was updated. + /// + /// Requires the server `vector` feature. + pub async fn vadd( + &mut self, + key: &str, + element: &str, + vector: &[f32], + ) -> Result { + let mut parts = Vec::with_capacity(3 + vector.len()); + parts.push(Frame::Bulk(Bytes::from_static(b"VADD"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(element.as_bytes()))); + for &f in vector { + parts.push(Frame::Bulk(Bytes::copy_from_slice( + f.to_string().as_bytes(), + ))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + match frame { + Frame::Integer(1) => Ok(true), + Frame::Integer(0) => Ok(false), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected 0 or 1 for VADD, got {other:?}" + ))), + } + } + + /// Adds multiple elements in a single round-trip. Returns the number of + /// newly inserted vectors (updates do not count). + /// + /// The server requires all vectors in a batch to have the same + /// dimensionality as the vector set (or creates it from the first batch). + pub async fn vadd_batch( + &mut self, + key: &str, + dim: usize, + entries: &[(&str, &[f32])], + ) -> Result { + let mut parts = Vec::with_capacity(5 + entries.len() * (1 + dim)); + parts.push(Frame::Bulk(Bytes::from_static(b"VADD_BATCH"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + parts.push(Frame::Bulk(Bytes::from_static(b"DIM"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice( + dim.to_string().as_bytes(), + ))); + for (element, vector) in entries { + parts.push(Frame::Bulk(Bytes::copy_from_slice(element.as_bytes()))); + for &f in *vector { + parts.push(Frame::Bulk(Bytes::copy_from_slice( + f.to_string().as_bytes(), + ))); + } + } + let frame = self.send_frame(Frame::Array(parts)).await?; + match frame { + Frame::Integer(n) => Ok(n), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected integer for VADD_BATCH, got {other:?}" + ))), + } + } + + /// Searches the vector set at `key` for the `count` nearest neighbours of + /// `query`. Returns results sorted by distance (closest first). + pub async fn vsim( + &mut self, + key: &str, + query: &[f32], + count: u32, + ) -> Result, ClientError> { + let mut parts = Vec::with_capacity(4 + query.len()); + parts.push(Frame::Bulk(Bytes::from_static(b"VSIM"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + for &f in query { + parts.push(Frame::Bulk(Bytes::copy_from_slice( + f.to_string().as_bytes(), + ))); + } + parts.push(Frame::Bulk(Bytes::from_static(b"COUNT"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice( + count.to_string().as_bytes(), + ))); + parts.push(Frame::Bulk(Bytes::from_static(b"WITHSCORES"))); + let frame = self.send_frame(Frame::Array(parts)).await?; + vsim_results(frame, true) + } + + /// Removes `element` from the vector set at `key`. Returns `true` if the + /// element was removed. + pub async fn vrem(&mut self, key: &str, element: &str) -> Result { + let frame = self + .send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"VREM")), + Frame::Bulk(Bytes::copy_from_slice(key.as_bytes())), + Frame::Bulk(Bytes::copy_from_slice(element.as_bytes())), + ])) + .await?; + match frame { + Frame::Integer(1) => Ok(true), + Frame::Integer(0) => Ok(false), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected 0 or 1 for VREM, got {other:?}" + ))), + } + } + + /// Returns the raw vector stored for `element` in the vector set at `key`, + /// or `None` if not found. + pub async fn vget( + &mut self, + key: &str, + element: &str, + ) -> Result>, ClientError> { + let frame = self + .send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"VGET")), + Frame::Bulk(Bytes::copy_from_slice(key.as_bytes())), + Frame::Bulk(Bytes::copy_from_slice(element.as_bytes())), + ])) + .await?; + vget_vector(frame) + } + + /// Returns the number of vectors in the vector set at `key`. + pub async fn vcard(&mut self, key: &str) -> Result { + let frame = self + .send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"VCARD")), + Frame::Bulk(Bytes::copy_from_slice(key.as_bytes())), + ])) + .await?; + match frame { + Frame::Integer(n) => Ok(n), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected integer for VCARD, got {other:?}" + ))), + } + } + + /// Returns the dimensionality of vectors in the set at `key`. + pub async fn vdim(&mut self, key: &str) -> Result { + let frame = self + .send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"VDIM")), + Frame::Bulk(Bytes::copy_from_slice(key.as_bytes())), + ])) + .await?; + match frame { + Frame::Integer(n) => Ok(n), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "expected integer for VDIM, got {other:?}" + ))), + } + } + + /// Returns metadata about the vector set at `key` as key-value pairs. + pub async fn vinfo(&mut self, key: &str) -> Result, ClientError> { + let frame = self + .send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"VINFO")), + Frame::Bulk(Bytes::copy_from_slice(key.as_bytes())), + ])) + .await?; + vinfo_pairs(frame) + } +} From b4259246f5804c113ee81c3a84ab2c3b5d6b2776 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Wed, 25 Feb 2026 13:03:32 -0500 Subject: [PATCH 5/5] style: cargo fmt cleanup --- crates/ember-client/src/commands.rs | 8 ++++---- tests/integration/src/client_typed_api.rs | 5 +---- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/crates/ember-client/src/commands.rs b/crates/ember-client/src/commands.rs index 2ab3a10f..a157d459 100644 --- a/crates/ember-client/src/commands.rs +++ b/crates/ember-client/src/commands.rs @@ -1076,9 +1076,7 @@ impl Client { /// Deletes keys asynchronously. Behaves like `del` but frees memory in /// the background for large values. Returns the number of keys removed. pub async fn unlink(&mut self, keys: &[&str]) -> Result { - let frame = self - .send_frame(cmd_keys_only(b"UNLINK", keys)) - .await?; + let frame = self.send_frame(cmd_keys_only(b"UNLINK", keys)).await?; integer(frame) } @@ -1106,7 +1104,9 @@ impl Client { /// message. pub async fn bgsave(&mut self) -> Result { let frame = self - .send_frame(Frame::Array(vec![Frame::Bulk(Bytes::from_static(b"BGSAVE"))])) + .send_frame(Frame::Array(vec![Frame::Bulk(Bytes::from_static( + b"BGSAVE", + ))])) .await?; string_value(frame) } diff --git a/tests/integration/src/client_typed_api.rs b/tests/integration/src/client_typed_api.rs index 5c42a44b..69278afb 100644 --- a/tests/integration/src/client_typed_api.rs +++ b/tests/integration/src/client_typed_api.rs @@ -423,10 +423,7 @@ async fn pubsub_channels_with_pattern_and_none() { #[tokio::test] async fn pubsub_numsub_empty_channels() { let (_server, mut client) = connect().await; - let pairs = client - .pubsub_numsub(&["no_one_here"]) - .await - .unwrap(); + let pairs = client.pubsub_numsub(&["no_one_here"]).await.unwrap(); assert_eq!(pairs.len(), 1); assert_eq!(pairs[0].0.as_ref(), b"no_one_here"); assert_eq!(pairs[0].1, 0);