BLPOP + COMMAND + MGET/MSET + DEL/EXISTS + EXPIRE/PEXPIRE/PERSIST + HINCRBY/HINCRBYFLOAT + BRPOP + DBSIZE + EXPIREAT/PEXIREAT implementations #1
| @@ -58,6 +58,7 @@ pub enum Cmd { | ||||
|     LPop(String, Option<u64>), | ||||
|     RPop(String, Option<u64>), | ||||
|     BLPop(Vec<String>, f64), | ||||
|     BRPop(Vec<String>, 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::<f64>() | ||||
|                                 .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<tokio::sync::oneshot::Receiver<(String, String)>> = 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<Protocol, DBError> { | ||||
|     // 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<u64> = Vec::with_capacity(keys.len()); | ||||
|     let mut names: Vec<String> = Vec::with_capacity(keys.len()); | ||||
|     let mut rxs: Vec<tokio::sync::oneshot::Receiver<(String, String)>> = 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<Protocol, DBError> { | ||||
|         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<Protocol, DBError> { | ||||
|     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![])), | ||||
|     } | ||||
| } | ||||
|   | ||||
| @@ -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; | ||||
|   | ||||
| @@ -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) | ||||
| } | ||||
|   | ||||
| @@ -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"); | ||||
| } | ||||
		Reference in New Issue
	
	Block a user