/build/source/nativelink-redis-tester/src/fake_redis.rs
Line | Count | Source |
1 | | // Copyright 2026 The NativeLink Authors. All rights reserved. |
2 | | // |
3 | | // Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); |
4 | | // you may not use this file except in compliance with the License. |
5 | | // You may obtain a copy of the License at |
6 | | // |
7 | | // See LICENSE file for details |
8 | | // |
9 | | // Unless required by applicable law or agreed to in writing, software |
10 | | // distributed under the License is distributed on an "AS IS" BASIS, |
11 | | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | | // See the License for the specific language governing permissions and |
13 | | // limitations under the License. |
14 | | |
15 | | use core::fmt::Write; |
16 | | use core::hash::BuildHasher; |
17 | | use std::collections::HashMap; |
18 | | |
19 | | use nativelink_util::background_spawn; |
20 | | use redis::Value; |
21 | | use redis_test::IntoRedisValue; |
22 | | use tokio::io::{AsyncReadExt, AsyncWriteExt}; |
23 | | use tokio::net::TcpListener; |
24 | | use tokio::sync::oneshot::{self, Sender}; |
25 | | use tracing::{error, info, warn}; |
26 | | |
27 | 109 | fn cmd_as_string(cmd: &redis::Cmd) -> String { |
28 | 109 | let raw = cmd.get_packed_command(); |
29 | 109 | String::from_utf8(raw).unwrap() |
30 | 109 | } |
31 | | |
32 | 992 | pub(crate) fn arg_as_string(output: &mut String, arg: Value) { |
33 | 992 | match arg { |
34 | 117 | Value::SimpleString(s) => { |
35 | 117 | write!(output, "+{s}\r\n").unwrap(); |
36 | 117 | } |
37 | 147 | Value::Okay => { |
38 | 147 | write!(output, "+OK\r\n").unwrap(); |
39 | 147 | } |
40 | 212 | Value::BulkString(s) => { |
41 | 212 | write!( |
42 | 212 | output, |
43 | 212 | "${}\r\n{}\r\n", |
44 | 212 | s.len(), |
45 | 212 | str::from_utf8(&s).unwrap() |
46 | 212 | ) |
47 | 212 | .unwrap(); |
48 | 212 | } |
49 | 206 | Value::Int(v) => { |
50 | 206 | write!(output, ":{v}\r\n").unwrap(); |
51 | 206 | } |
52 | 192 | Value::Array(values) => { |
53 | 192 | write!(output, "*{}\r\n", values.len()).unwrap(); |
54 | 376 | for value in values192 { |
55 | 376 | arg_as_string(output, value); |
56 | 376 | } |
57 | | } |
58 | 65 | Value::Map(values) => { |
59 | 65 | write!(output, "%{}\r\n", values.len()).unwrap(); |
60 | 86 | for (key, value) in values65 { |
61 | 86 | arg_as_string(output, key); |
62 | 86 | arg_as_string(output, value); |
63 | 86 | } |
64 | | } |
65 | 52 | Value::Nil => { |
66 | 52 | write!(output, "_\r\n").unwrap(); |
67 | 52 | } |
68 | 1 | Value::Boolean(value) => { |
69 | 1 | if value { |
70 | 1 | write!(output, "#t\r\n") |
71 | | } else { |
72 | 0 | write!(output, "#f\r\n") |
73 | | } |
74 | 1 | .unwrap(); |
75 | | } |
76 | | _ => { |
77 | 0 | panic!("No support for {arg:?}") |
78 | | } |
79 | | } |
80 | 992 | } |
81 | | |
82 | 89 | fn args_as_string(args: Vec<Value>) -> String { |
83 | 89 | let mut output = String::new(); |
84 | 215 | for arg in args89 { |
85 | 215 | arg_as_string(&mut output, arg); |
86 | 215 | } |
87 | 89 | output |
88 | 89 | } |
89 | | |
90 | 49 | pub fn add_to_response<B: BuildHasher>( |
91 | 49 | response: &mut HashMap<String, String, B>, |
92 | 49 | cmd: &redis::Cmd, |
93 | 49 | args: Vec<Value>, |
94 | 49 | ) { |
95 | 49 | add_to_response_raw(response, cmd, args_as_string(args)); |
96 | 49 | } |
97 | | |
98 | 49 | pub fn add_to_response_raw<B: BuildHasher>( |
99 | 49 | response: &mut HashMap<String, String, B>, |
100 | 49 | cmd: &redis::Cmd, |
101 | 49 | args: String, |
102 | 49 | ) { |
103 | 49 | response.insert(cmd_as_string(cmd), args); |
104 | 49 | } |
105 | | |
106 | 20 | fn setinfo(responses: &mut HashMap<String, String>) { |
107 | | // We do raw inserts of command here, because the library sends 3/4 commands in one go |
108 | | // They always start with HELLO, then optionally SELECT, so we use this to differentiate |
109 | 20 | let hello = cmd_as_string(redis::cmd("HELLO").arg("3")); |
110 | 20 | let setinfo = cmd_as_string( |
111 | 20 | redis::cmd("CLIENT") |
112 | 20 | .arg("SETINFO") |
113 | 20 | .arg("LIB-NAME") |
114 | 20 | .arg("redis-rs"), |
115 | | ); |
116 | 20 | responses.insert( |
117 | 20 | [hello.clone(), setinfo.clone()].join(""), |
118 | 20 | args_as_string(vec![ |
119 | 20 | Value::Map(vec![( |
120 | 20 | Value::SimpleString("server".into()), |
121 | 20 | Value::SimpleString("redis".into()), |
122 | 20 | )]), |
123 | 20 | Value::Okay, |
124 | 20 | Value::Okay, |
125 | | ]), |
126 | | ); |
127 | 20 | responses.insert( |
128 | 20 | [hello, cmd_as_string(redis::cmd("SELECT").arg(3)), setinfo].join(""), |
129 | 20 | args_as_string(vec![ |
130 | 20 | Value::Map(vec![( |
131 | 20 | Value::SimpleString("server".into()), |
132 | 20 | Value::SimpleString("redis".into()), |
133 | 20 | )]), |
134 | 20 | Value::Okay, |
135 | 20 | Value::Okay, |
136 | 20 | Value::Okay, |
137 | | ]), |
138 | | ); |
139 | 20 | } |
140 | | |
141 | 13 | pub fn add_lua_script<B: BuildHasher>( |
142 | 13 | responses: &mut HashMap<String, String, B>, |
143 | 13 | lua_script: &str, |
144 | 13 | hash: &str, |
145 | 13 | ) { |
146 | 13 | add_to_response( |
147 | 13 | responses, |
148 | 13 | redis::cmd("SCRIPT").arg("LOAD").arg(lua_script), |
149 | 13 | vec![hash.into_redis_value()], |
150 | | ); |
151 | 13 | } |
152 | | |
153 | 13 | pub fn fake_redis_stream() -> HashMap<String, String> { |
154 | 13 | let mut responses = HashMap::new(); |
155 | 13 | setinfo(&mut responses); |
156 | | // Does setinfo as well, so need to respond to all 3 |
157 | 13 | add_to_response( |
158 | 13 | &mut responses, |
159 | 13 | redis::cmd("SELECT").arg("3"), |
160 | 13 | vec![Value::Okay, Value::Okay, Value::Okay], |
161 | | ); |
162 | 13 | responses |
163 | 13 | } |
164 | | |
165 | 3 | pub fn fake_redis_sentinel_master_stream() -> HashMap<String, String> { |
166 | 3 | let mut response = fake_redis_stream(); |
167 | 3 | add_to_response( |
168 | 3 | &mut response, |
169 | 3 | &redis::cmd("ROLE"), |
170 | 3 | vec![Value::Array(vec![ |
171 | 3 | "master".into_redis_value(), |
172 | 3 | 0.into_redis_value(), |
173 | 3 | Value::Array(vec![]), |
174 | 3 | ])], |
175 | | ); |
176 | 3 | response |
177 | 3 | } |
178 | | |
179 | 7 | pub fn fake_redis_sentinel_stream(master_name: &str, redis_port: u16) -> HashMap<String, String> { |
180 | 7 | let mut response = HashMap::new(); |
181 | 7 | setinfo(&mut response); |
182 | | |
183 | | // Not a full "sentinel masters" response, but enough for redis-rs |
184 | 7 | let resp: Vec<(Value, Value)> = vec![ |
185 | 7 | ("name".into_redis_value(), master_name.into_redis_value()), |
186 | 7 | ("ip".into_redis_value(), "127.0.0.1".into_redis_value()), |
187 | 7 | ( |
188 | 7 | "port".into_redis_value(), |
189 | 7 | i64::from(redis_port).into_redis_value(), |
190 | 7 | ), |
191 | 7 | ("flags".into_redis_value(), "master".into_redis_value()), |
192 | | ]; |
193 | | |
194 | 7 | add_to_response( |
195 | 7 | &mut response, |
196 | 7 | redis::cmd("SENTINEL").arg("MASTERS"), |
197 | 7 | vec![Value::Array(vec![Value::Map(resp)])], |
198 | | ); |
199 | 7 | response |
200 | 7 | } |
201 | | |
202 | 28 | pub(crate) async fn fake_redis_internal<H>( |
203 | 28 | listener: TcpListener, |
204 | 28 | listener_ready_tx: Sender<()>, |
205 | 28 | handlers: Vec<H>, |
206 | 28 | ) where |
207 | 28 | H: Fn(&[u8]) -> String + Send + Clone + 'static + Sync, |
208 | 28 | { |
209 | 28 | let mut handler_iter = handlers.iter().cloned().cycle(); |
210 | 28 | info!( |
211 | | "Waiting for connection on {}", |
212 | 28 | listener.local_addr().unwrap() |
213 | | ); |
214 | 28 | listener_ready_tx |
215 | 28 | .send(()) |
216 | 28 | .expect("Expected successful send"); |
217 | | loop { |
218 | 70 | let Ok((mut stream42 , _)) = listener.accept().await else { |
219 | 0 | error!("accept error"); |
220 | 0 | panic!("error"); |
221 | | }; |
222 | 42 | info!("Accepted new connection"); |
223 | 42 | let local_handler = handler_iter.next().unwrap(); |
224 | 42 | background_spawn!("thread", async move { |
225 | | loop { |
226 | 103k | let mut buf = vec![0; 8192]; |
227 | 103k | let res103k = stream.read(&mut buf).await.unwrap103k (); |
228 | 103k | if res != 0 { |
229 | 526 | let output = local_handler(&buf[..res]); |
230 | 526 | if !output.is_empty() { |
231 | 262 | stream.write_all(output.as_bytes()).await.unwrap(); |
232 | 264 | } |
233 | 103k | } |
234 | | } |
235 | | }); |
236 | | } |
237 | | } |
238 | | |
239 | 21 | async fn fake_redis<B>( |
240 | 21 | listener: TcpListener, |
241 | 21 | listener_ready_tx: Sender<()>, |
242 | 21 | all_responses: Vec<HashMap<String, String, B>>, |
243 | 21 | ) where |
244 | 21 | B: BuildHasher + Clone + Send + 'static + Sync, |
245 | 21 | { |
246 | 21 | let funcs = all_responses |
247 | 21 | .iter() |
248 | 21 | .map(|responses| { |
249 | 21 | info!("Responses are: {:?}", responses); |
250 | 21 | let values = responses.clone(); |
251 | 329 | move |buf: &[u8]| -> String { |
252 | 329 | let str_buf = String::from_utf8_lossy(buf).into_owned(); |
253 | 1.24k | for (key, value) in &values329 { |
254 | 1.24k | if str_buf.starts_with(key) { |
255 | 65 | info!("Responding to {}", str_buf.replace("\r\n", "\\r\\n")); |
256 | 65 | return value.clone(); |
257 | 1.17k | } |
258 | | } |
259 | 264 | warn!( |
260 | | "Unknown command: {}", |
261 | 264 | str_buf.chars().take(1000).collect::<String>() |
262 | | ); |
263 | 264 | String::new() |
264 | 329 | } |
265 | 21 | }) |
266 | 21 | .collect(); |
267 | 21 | fake_redis_internal(listener, listener_ready_tx, funcs).await; |
268 | 0 | } |
269 | | |
270 | 21 | async fn make_fake_redis_with_multiple_responses<B: BuildHasher + Clone + Send + 'static + Sync>( |
271 | 21 | responses: Vec<HashMap<String, String, B>>, |
272 | 21 | ) -> u16 { |
273 | 21 | let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); |
274 | 21 | let port = listener.local_addr().unwrap().port(); |
275 | | |
276 | 21 | let (listener_ready_tx, listener_ready_rx) = oneshot::channel::<()>(); |
277 | | |
278 | 21 | background_spawn!("listener", async move { |
279 | 21 | fake_redis(listener, listener_ready_tx, responses).await; |
280 | 0 | }); |
281 | | |
282 | 21 | listener_ready_rx |
283 | 21 | .await |
284 | 21 | .expect("Expected successful listener boot"); |
285 | 21 | info!(port, "Fake redis booted"); |
286 | | |
287 | 21 | port |
288 | 21 | } |
289 | | |
290 | 21 | pub async fn make_fake_redis_with_responses<B: BuildHasher + Clone + Send + 'static + Sync>( |
291 | 21 | responses: HashMap<String, String, B>, |
292 | 21 | ) -> u16 { |
293 | 21 | make_fake_redis_with_multiple_responses(vec![responses]).await |
294 | 21 | } |