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("timeout Some(0) is ambiguous (wire 0 = wait forever); use None to wait forever".into()));
88 }
89 Ok(())
90}
91
92fn blocking_request(
96 c: &mut RespClient,
97 verb: &[u8],
98 keys: &[&[u8]],
99 timeout: Option<Duration>,
100) -> KevyResult<Reply> {
101 let mut args = Vec::with_capacity(keys.len() + 2);
102 args.push(verb.to_vec());
103 args.extend(keys.iter().map(|k| k.to_vec()));
104 args.push(match timeout {
105 None => b"0".to_vec(),
106 Some(d) => format!("{}", d.as_secs_f64()).into_bytes(),
107 });
108 Ok(c.request(&args)?)
109}
110
111fn pop_kv(reply: Reply) -> KevyResult<Option<(Vec<u8>, Vec<u8>)>> {
113 match reply {
114 Reply::Array(items) if items.len() == 2 => {
115 let mut it = items.into_iter();
116 match (it.next().unwrap(), it.next().unwrap()) {
117 (Reply::Bulk(k), Reply::Bulk(v)) => Ok(Some((k, v))),
118 (a, _) => Err(unexpected(a)),
119 }
120 }
121 Reply::Nil | Reply::Null => Ok(None),
122 Reply::Error(e) => Err(KevyError::Protocol(string(e))),
123 other => Err(unexpected(other)),
124 }
125}
126
127fn pop_kv_score(reply: Reply) -> KevyResult<Option<ZPopHit>> {
129 match reply {
130 Reply::Array(items) if items.len() == 3 => {
131 let mut it = items.into_iter();
132 match (it.next().unwrap(), it.next().unwrap(), it.next().unwrap()) {
133 (Reply::Bulk(k), Reply::Bulk(m), Reply::Bulk(s)) => {
134 Ok(Some((k, m, num_f64(&s)?)))
135 }
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
163 .blpop(&[&b"empty"[..]], Some(Duration::from_millis(30)))
164 .unwrap();
165 assert_eq!(hit, None);
166 }
167
168 #[test]
169 fn embedded_bzpopmin_pops_lowest_score() {
170 let mut c = Connection::connect("mem://").unwrap();
171 c.zadd(b"z", &[(2.0, b"hi".as_ref()), (1.0, b"lo".as_ref())])
172 .unwrap();
173 let hit = c
174 .bzpopmin(&[&b"z"[..]], Some(Duration::from_secs(1)))
175 .unwrap();
176 assert_eq!(hit, Some((b"z".to_vec(), b"lo".to_vec(), 1.0)));
177 let miss = c
178 .bzpopmin(&[&b"gone"[..]], Some(Duration::from_millis(30)))
179 .unwrap();
180 assert_eq!(miss, None);
181 }
182
183 #[test]
184 fn zero_some_timeout_and_empty_keys_rejected() {
185 let mut c = Connection::connect("mem://").unwrap();
186 let err = c.blpop(&[&b"q"[..]], Some(Duration::ZERO)).unwrap_err();
187 assert!(matches!(err, KevyError::InvalidInput(_)));
188 let err = c.blpop(&[], Some(Duration::from_secs(1))).unwrap_err();
189 assert!(matches!(err, KevyError::InvalidInput(_)));
190 }
191}