1use crate::{KevyError, KevyResult};
17use std::time::Duration;
18
19use kevy_resp::Reply;
20use kevy_resp_client::RespClient;
21
22use crate::{Connection, num_f64, string, unexpected};
23
24pub type ZPopHit = (Vec<u8>, Vec<u8>, f64);
26
27impl Connection {
28 pub fn blpop(
36 &mut self,
37 keys: &[&[u8]],
38 timeout: Option<Duration>,
39 ) -> KevyResult<Option<(Vec<u8>, Vec<u8>)>> {
40 check_blocking_args(keys, timeout)?;
41 match self {
42 Self::Embedded(s) => s.blpop(keys, timeout),
43 Self::Remote(c) => pop_kv(blocking_request(c, b"BLPOP", keys, timeout)?),
44 }
45 }
46
47 pub fn brpop(
50 &mut self,
51 keys: &[&[u8]],
52 timeout: Option<Duration>,
53 ) -> KevyResult<Option<(Vec<u8>, Vec<u8>)>> {
54 check_blocking_args(keys, timeout)?;
55 match self {
56 Self::Embedded(s) => s.brpop(keys, timeout),
57 Self::Remote(c) => pop_kv(blocking_request(c, b"BRPOP", keys, timeout)?),
58 }
59 }
60
61 pub fn bzpopmin(
67 &mut self,
68 keys: &[&[u8]],
69 timeout: Option<Duration>,
70 ) -> KevyResult<Option<ZPopHit>> {
71 check_blocking_args(keys, timeout)?;
72 match self {
73 Self::Embedded(s) => s.bzpopmin(keys, timeout),
74 Self::Remote(c) => pop_kv_score(blocking_request(c, b"BZPOPMIN", keys, timeout)?),
75 }
76 }
77}
78
79fn check_blocking_args(keys: &[&[u8]], timeout: Option<Duration>) -> KevyResult<()> {
83 if keys.is_empty() {
84 return Err(KevyError::InvalidInput("blocking pop needs at least one key".into()));
85 }
86 if timeout == Some(Duration::ZERO) {
87 return Err(KevyError::InvalidInput(
88 "timeout Some(0) is ambiguous (wire 0 = wait forever); use None to wait forever".into(),
89 ));
90 }
91 Ok(())
92}
93
94fn blocking_request(
98 c: &mut RespClient,
99 verb: &[u8],
100 keys: &[&[u8]],
101 timeout: Option<Duration>,
102) -> KevyResult<Reply> {
103 let mut args = Vec::with_capacity(keys.len() + 2);
104 args.push(verb.to_vec());
105 args.extend(keys.iter().map(|k| k.to_vec()));
106 args.push(match timeout {
107 None => b"0".to_vec(),
108 Some(d) => format!("{}", d.as_secs_f64()).into_bytes(),
109 });
110 Ok(c.request(&args)?)
111}
112
113fn pop_kv(reply: Reply) -> KevyResult<Option<(Vec<u8>, Vec<u8>)>> {
115 match reply {
116 Reply::Array(items) if items.len() == 2 => {
117 let mut it = items.into_iter();
118 match (it.next().unwrap(), it.next().unwrap()) {
119 (Reply::Bulk(k), Reply::Bulk(v)) => Ok(Some((k, v))),
120 (a, _) => Err(unexpected(a)),
121 }
122 }
123 Reply::Nil | Reply::Null => Ok(None),
124 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
125 other => Err(unexpected(other)),
126 }
127}
128
129fn pop_kv_score(reply: Reply) -> KevyResult<Option<ZPopHit>> {
131 match reply {
132 Reply::Array(items) if items.len() == 3 => {
133 let mut it = items.into_iter();
134 match (it.next().unwrap(), it.next().unwrap(), it.next().unwrap()) {
135 (Reply::Bulk(k), Reply::Bulk(m), Reply::Bulk(s)) => Ok(Some((k, m, num_f64(&s)?))),
136 (a, _, _) => Err(unexpected(a)),
137 }
138 }
139 Reply::Nil | Reply::Null => Ok(None),
140 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
141 other => Err(unexpected(other)),
142 }
143}
144
145#[cfg(test)]
146mod tests {
147 use super::*;
148
149 #[test]
150 fn embedded_blpop_immediate_hit() {
151 let mut c = Connection::connect("mem://").unwrap();
152 c.rpush(b"q", &[&b"a"[..], &b"b"[..]]).unwrap();
153 let hit = c.blpop(&[&b"q"[..]], Some(Duration::from_secs(1))).unwrap();
154 assert_eq!(hit, Some((b"q".to_vec(), b"a".to_vec())));
155 let hit = c.brpop(&[&b"q"[..]], Some(Duration::from_secs(1))).unwrap();
156 assert_eq!(hit, Some((b"q".to_vec(), b"b".to_vec())));
157 }
158
159 #[test]
160 fn embedded_blpop_timeout_returns_none() {
161 let mut c = Connection::connect("mem://").unwrap();
162 let hit = c.blpop(&[&b"empty"[..]], Some(Duration::from_millis(30))).unwrap();
163 assert_eq!(hit, None);
164 }
165
166 #[test]
167 fn embedded_bzpopmin_pops_lowest_score() {
168 let mut c = Connection::connect("mem://").unwrap();
169 c.zadd(b"z", &[(2.0, b"hi".as_ref()), (1.0, b"lo".as_ref())]).unwrap();
170 let hit = c.bzpopmin(&[&b"z"[..]], Some(Duration::from_secs(1))).unwrap();
171 assert_eq!(hit, Some((b"z".to_vec(), b"lo".to_vec(), 1.0)));
172 let miss = c.bzpopmin(&[&b"gone"[..]], Some(Duration::from_millis(30))).unwrap();
173 assert_eq!(miss, None);
174 }
175
176 #[test]
177 fn zero_some_timeout_and_empty_keys_rejected() {
178 let mut c = Connection::connect("mem://").unwrap();
179 let err = c.blpop(&[&b"q"[..]], Some(Duration::ZERO)).unwrap_err();
180 assert!(matches!(err, KevyError::InvalidInput(_)));
181 let err = c.blpop(&[], Some(Duration::from_secs(1))).unwrap_err();
182 assert!(matches!(err, KevyError::InvalidInput(_)));
183 }
184}