From 463000c8f7ba3e9e29ac517c4b509b142be75c12 Mon Sep 17 00:00:00 2001 From: Maxime Van Hees Date: Tue, 19 Aug 2025 15:52:36 +0200 Subject: [PATCH] Implemented BRPOP and minimal COMMAND DOCS stub, and wired side-aware waiter delivery --- herodb/src/cmd.rs | 115 +++++++++++++++++++++++- herodb/src/server.rs | 18 +++- herodb/tests/redis_integration_tests.rs | 8 +- herodb/tests/usage_suite.rs | 36 ++++++++ 4 files changed, 167 insertions(+), 10 deletions(-) diff --git a/herodb/src/cmd.rs b/herodb/src/cmd.rs index 8a9058c..d4287ef 100644 --- a/herodb/src/cmd.rs +++ b/herodb/src/cmd.rs @@ -58,6 +58,7 @@ pub enum Cmd { LPop(String, Option), RPop(String, Option), BLPop(Vec, f64), + BRPop(Vec, f64), LLen(String), LRem(String, i64, String), LTrim(String, i64, i64), @@ -503,6 +504,17 @@ impl Cmd { .map_err(|_| DBError("ERR timeout is not a number".to_string()))?; Cmd::BLPop(keys, timeout_f) } + "brpop" => { + if cmd.len() < 3 { + return Err(DBError(format!("wrong number of arguments for BRPOP command"))); + } + // keys are all but the last argument + let keys = cmd[1..cmd.len()-1].to_vec(); + let timeout_f = cmd[cmd.len()-1] + .parse::() + .map_err(|_| DBError("ERR timeout is not a number".to_string()))?; + Cmd::BRPop(keys, timeout_f) + } "llen" => { if cmd.len() != 2 { return Err(DBError(format!("wrong number of arguments for LLEN command"))); @@ -663,13 +675,14 @@ impl Cmd { Cmd::Client(_) => Ok(Protocol::SimpleString("OK".to_string())), Cmd::ClientSetName(name) => client_setname_cmd(server, &name).await, Cmd::ClientGetName => client_getname_cmd(server).await, - Cmd::Command(_) => Ok(Protocol::Array(vec![])), + Cmd::Command(args) => command_cmd(&args), // List commands Cmd::LPush(key, elements) => lpush_cmd(server, &key, &elements).await, Cmd::RPush(key, elements) => rpush_cmd(server, &key, &elements).await, Cmd::LPop(key, count) => lpop_cmd(server, &key, &count).await, Cmd::RPop(key, count) => rpop_cmd(server, &key, &count).await, Cmd::BLPop(keys, timeout) => blpop_cmd(server, &keys, timeout).await, + Cmd::BRPop(keys, timeout) => brpop_cmd(server, &keys, timeout).await, Cmd::LLen(key) => llen_cmd(server, &key).await, Cmd::LRem(key, count, element) => lrem_cmd(server, &key, count, &element).await, Cmd::LTrim(key, start, stop) => ltrim_cmd(server, &key, start, stop).await, @@ -825,7 +838,89 @@ async fn blpop_cmd(server: &Server, keys: &[String], timeout_secs: f64) -> Resul let mut rxs: Vec> = Vec::with_capacity(keys.len()); for k in keys { - let (id, rx) = server.register_waiter(db_index, k).await; + let (id, rx) = server.register_waiter(db_index, k, crate::server::PopSide::Left).await; + ids.push(id); + names.push(k.clone()); + rxs.push(rx); + } + + // Wait for the first delivery or timeout + let wait_fut = async move { + let mut futures_vec = rxs; + loop { + if futures_vec.is_empty() { + return None; + } + let (res, idx, remaining) = select_all(futures_vec).await; + match res { + Ok((k, elem)) => { + return Some((k, elem, idx, remaining)); + } + Err(_canceled) => { + // That waiter was canceled; continue with the rest + futures_vec = remaining; + continue; + } + } + } + }; + + match timeout(Duration::from_secs_f64(timeout_secs), wait_fut).await { + Ok(Some((k, elem, idx, _remaining))) => { + // Unregister other waiters + for (i, key_name) in names.iter().enumerate() { + if i != idx { + server.unregister_waiter(db_index, key_name, ids[i]).await; + } + } + Ok(Protocol::Array(vec![ + Protocol::BulkString(k), + Protocol::BulkString(elem), + ])) + } + Ok(None) => { + // No futures left; unregister all waiters + for (i, key_name) in names.iter().enumerate() { + server.unregister_waiter(db_index, key_name, ids[i]).await; + } + Ok(Protocol::Null) + } + Err(_elapsed) => { + // Timeout: unregister all waiters + for (i, key_name) in names.iter().enumerate() { + server.unregister_waiter(db_index, key_name, ids[i]).await; + } + Ok(Protocol::Null) + } + } +} + +// BRPOP implementation (mirror of BLPOP, popping from the right) +async fn brpop_cmd(server: &Server, keys: &[String], timeout_secs: f64) -> Result { + // Immediate, non-blocking attempt in key order using RPOP + for k in keys { + let elems = server.current_storage()?.rpop(k, 1)?; + if !elems.is_empty() { + return Ok(Protocol::Array(vec![ + Protocol::BulkString(k.clone()), + Protocol::BulkString(elems[0].clone()), + ])); + } + } + + // If timeout is zero, return immediately with Null + if timeout_secs <= 0.0 { + return Ok(Protocol::Null); + } + + // Register waiters for each key (Right side) + let db_index = server.selected_db; + let mut ids: Vec = Vec::with_capacity(keys.len()); + let mut names: Vec = Vec::with_capacity(keys.len()); + let mut rxs: Vec> = Vec::with_capacity(keys.len()); + + for k in keys { + let (id, rx) = server.register_waiter(db_index, k, crate::server::PopSide::Right).await; ids.push(id); names.push(k.clone()); rxs.push(rx); @@ -1358,3 +1453,19 @@ async fn client_getname_cmd(server: &Server) -> Result { None => Ok(Protocol::Null), } } + +// Minimal COMMAND subcommands stub to satisfy redis-cli probes. +// - COMMAND DOCS ... => return empty array +// - COMMAND INFO ... => return empty array +// - Any other => empty array +fn command_cmd(args: &[String]) -> Result { + if args.is_empty() { + return Ok(Protocol::Array(vec![])); + } + let sub = args[0].to_lowercase(); + match sub.as_str() { + "docs" => Ok(Protocol::Array(vec![])), + "info" => Ok(Protocol::Array(vec![])), + _ => Ok(Protocol::Array(vec![])), + } +} diff --git a/herodb/src/server.rs b/herodb/src/server.rs index 23a93af..0c128c6 100644 --- a/herodb/src/server.rs +++ b/herodb/src/server.rs @@ -28,9 +28,16 @@ pub struct Server { pub struct Waiter { pub id: u64, + pub side: PopSide, pub tx: oneshot::Sender<(String, String)>, // (key, element) } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PopSide { + Left, + Right, +} + impl Server { pub async fn new(option: options::DBOption) -> Self { Server { @@ -83,14 +90,14 @@ impl Server { // ----- BLPOP waiter helpers ----- - pub async fn register_waiter(&self, db_index: u64, key: &str) -> (u64, oneshot::Receiver<(String, String)>) { + pub async fn register_waiter(&self, db_index: u64, key: &str, side: PopSide) -> (u64, oneshot::Receiver<(String, String)>) { let id = self.waiter_seq.fetch_add(1, Ordering::Relaxed); let (tx, rx) = oneshot::channel::<(String, String)>(); let mut guard = self.list_waiters.lock().await; let per_db = guard.entry(db_index).or_insert_with(HashMap::new); let q = per_db.entry(key.to_string()).or_insert_with(Vec::new); - q.push(Waiter { id, tx }); + q.push(Waiter { id, side, tx }); (id, rx) } @@ -135,8 +142,11 @@ impl Server { let waiter = if let Some(w) = maybe_waiter { w } else { break }; - // Pop one element from the left - let elems = self.current_storage()?.lpop(key, 1)?; + // Pop one element depending on waiter side + let elems = match waiter.side { + PopSide::Left => self.current_storage()?.lpop(key, 1)?, + PopSide::Right => self.current_storage()?.rpop(key, 1)?, + }; if elems.is_empty() { // Nothing to deliver; re-register waiter at the front to preserve order let mut guard = self.list_waiters.lock().await; diff --git a/herodb/tests/redis_integration_tests.rs b/herodb/tests/redis_integration_tests.rs index 16e1f64..47033e1 100644 --- a/herodb/tests/redis_integration_tests.rs +++ b/herodb/tests/redis_integration_tests.rs @@ -16,9 +16,9 @@ fn get_redis_connection(port: u16) -> Connection { } } Err(e) => { - if attempts >= 20 { + if attempts >= 120 { panic!( - "Failed to connect to Redis server after 20 attempts: {}", + "Failed to connect to Redis server after 120 attempts: {}", e ); } @@ -88,8 +88,8 @@ fn setup_server() -> (ServerProcessGuard, u16) { test_dir, }; - // Give the server a moment to start - std::thread::sleep(Duration::from_millis(500)); + // Give the server time to build and start (cargo run may compile first) + std::thread::sleep(Duration::from_millis(2500)); (guard, port) } diff --git a/herodb/tests/usage_suite.rs b/herodb/tests/usage_suite.rs index c61fecf..9591193 100644 --- a/herodb/tests/usage_suite.rs +++ b/herodb/tests/usage_suite.rs @@ -779,4 +779,40 @@ async fn test_12_hash_incr() { let _ = send_cmd(&mut s, &["HSET", "hinc", "notf", "abc"]).await; let r6 = send_cmd(&mut s, &["HINCRBYFLOAT", "hinc", "notf", "1"]).await; assert_contains(&r6, "ERR", "HINCRBYFLOAT on non-float field should ERR"); +} +#[tokio::test] +async fn test_05b_brpop_suite() { + let (server, port) = start_test_server("lists_brpop").await; + spawn_listener(server, port).await; + sleep(Duration::from_millis(150)).await; + + let mut a = connect(port).await; + + // RPUSH some initial data, BRPOP should take from the right + let _ = send_cmd(&mut a, &["RPUSH", "q:rjobs", "1", "2"]).await; + let br_nonblock = send_cmd(&mut a, &["BRPOP", "q:rjobs", "0"]).await; + // Should pop the rightmost element "2" + assert_contains(&br_nonblock, "q:rjobs", "BRPOP returns key"); + assert_contains(&br_nonblock, "2", "BRPOP returns rightmost element"); + + // Now test blocking BRPOP: start blocked client, then RPUSH from another client + let c1 = connect(port).await; + let mut c2 = connect(port).await; + + // Start BRPOP on c1 + let brpop_task = tokio::spawn(async move { + let mut c1_local = c1; + send_cmd(&mut c1_local, &["BRPOP", "q:blockr", "5"]).await + }); + + // Give it time to register waiter + sleep(Duration::from_millis(150)).await; + + // Push from right to wake BRPOP + let _ = send_cmd(&mut c2, &["RPUSH", "q:blockr", "X"]).await; + + // Await BRPOP result + let brpop_res = brpop_task.await.expect("BRPOP task join"); + assert_contains(&brpop_res, "q:blockr", "BRPOP returned key"); + assert_contains(&brpop_res, "X", "BRPOP returned element"); } \ No newline at end of file