/build/source/nativelink-redis-tester/src/dynamic_fake_redis.rs
Line | Count | Source |
1 | | // Copyright 2026 Trace Machina, Inc. 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; |
16 | | use core::time::Duration; |
17 | | use std::collections::HashMap; |
18 | | use std::collections::hash_map::Entry; |
19 | | use std::sync::{Arc, Mutex}; |
20 | | |
21 | | use nativelink_util::background_spawn; |
22 | | use redis::Value; |
23 | | use redis_protocol::resp2::decode::decode; |
24 | | use redis_protocol::resp2::types::{OwnedFrame, Resp2Frame}; |
25 | | use tokio::io::{AsyncReadExt, AsyncWriteExt}; |
26 | | use tokio::net::TcpListener; |
27 | | use tokio::sync::oneshot::{self, Sender}; |
28 | | use tracing::{debug, info, trace}; |
29 | | |
30 | | use crate::fake_redis::arg_as_string; |
31 | | |
32 | | /// The error a command answers with instead of running. |
33 | | #[derive(Clone, Debug, PartialEq, Eq)] |
34 | | pub enum Fault { |
35 | | /// `-ERR <message>`. |
36 | | Error(String), |
37 | | /// `-READONLY`, what a replica answers a write with after a failover. |
38 | | ReadOnly, |
39 | | /// `-OOM`, what a full server with `noeviction` answers a write with. |
40 | | Oom, |
41 | | /// The command runs, then the client reads `-ERR <message>`: a reply |
42 | | /// lost after the write landed, as a timeout loses it. |
43 | | AppliedThenError(String), |
44 | | } |
45 | | |
46 | | impl Fault { |
47 | 36 | fn reply(&self) -> String { |
48 | 36 | match self { |
49 | 3 | Self::Error(message) | Self::AppliedThenError(message1 ) => { |
50 | 4 | format!("-ERR {message}\r\n") |
51 | | } |
52 | 28 | Self::ReadOnly => "-READONLY You can't write against a read only replica.\r\n".into(), |
53 | 4 | Self::Oom => "-OOM command not allowed when used memory > 'maxmemory'.\r\n".into(), |
54 | | } |
55 | 36 | } |
56 | | } |
57 | | |
58 | | #[derive(Clone, Debug)] |
59 | | struct FaultRule { |
60 | | /// Upper-case command name, or `*` for any. |
61 | | command: String, |
62 | | /// Matches commands whose first argument starts with this. |
63 | | key_prefix: Option<String>, |
64 | | fault: Fault, |
65 | | /// How many more commands it fails; `None` is until cleared. |
66 | | remaining: Option<u64>, |
67 | | } |
68 | | |
69 | | /// What the fake does wrong, consulted on every command and connection. |
70 | | /// The default does nothing wrong. |
71 | | #[derive(Debug, Default)] |
72 | | struct FaultPolicy { |
73 | | rules: Vec<FaultRule>, |
74 | | /// Commands received, by upper-case name, faulted ones included. |
75 | | counts: HashMap<String, u64>, |
76 | | /// Commands received since the start; never reset. |
77 | | total: u64, |
78 | | /// Connections are closed and new ones refused while set. |
79 | | partitioned: bool, |
80 | | /// Bumped to close every open connection once. |
81 | | generation: u64, |
82 | | /// Held before each reply is written. |
83 | | latency: Option<Duration>, |
84 | | /// Ports whose listener is gone: its connections close and it accepts |
85 | | /// no more, as when the process that held them dies. |
86 | | retired: std::collections::HashSet<u16>, |
87 | | } |
88 | | |
89 | | impl FaultPolicy { |
90 | | /// The reply that replaces the command's, and whether the command |
91 | | /// still runs. |
92 | 761k | fn on_command(&mut self, command: &str, key: Option<&str>) -> Option<(String, bool)> { |
93 | 761k | *self.counts.entry(command.to_string()).or_default() += 1; |
94 | 761k | self.total += 1; |
95 | 761k | let index36 = self.rules.iter().position(|rule| {486 |
96 | 486 | (rule.command == "*" || rule.command == command) |
97 | 36 | && rule |
98 | 36 | .key_prefix |
99 | 36 | .as_ref() |
100 | 36 | .is_none_or(|prefix| key0 .is_some_and0 (|key| key0 .starts_with0 (prefix0 .as_str0 ()))) |
101 | 761k | }486 )?; |
102 | 36 | let rule = &mut self.rules[index]; |
103 | 36 | let reply = ( |
104 | 36 | rule.fault.reply(), |
105 | 36 | matches!35 (rule.fault, Fault::AppliedThenError(_)), |
106 | | ); |
107 | 36 | if let Some(remaining33 ) = &mut rule.remaining { |
108 | 33 | *remaining -= 1; |
109 | 33 | if *remaining == 0 { |
110 | 8 | self.rules.remove(index); |
111 | 25 | } |
112 | 3 | } |
113 | 36 | Some(reply) |
114 | 761k | } |
115 | | } |
116 | | |
117 | | /// How many leading bytes of `buf` are whole RESP frames. |
118 | 522k | fn complete_frames_len(buf: &[u8]) -> usize { |
119 | 522k | let mut index = 0; |
120 | 1.28M | while index < buf.len() { |
121 | 762k | match decode(&buf[index..]) { |
122 | 761k | Ok(Some((_, amt))) => index += amt, |
123 | 539 | _ => break, |
124 | | } |
125 | | } |
126 | 522k | index |
127 | 522k | } |
128 | | |
129 | | /// Commands that write, for faults aimed at writes only. |
130 | | pub const WRITE_COMMANDS: &[&str] = &["EVALSHA", "HMSET", "EXPIRE", "PUBLISH"]; |
131 | | |
132 | | pub trait SubscriptionManagerNotify { |
133 | | fn notify_for_test(&self, value: String); |
134 | | } |
135 | | |
136 | | #[derive(Clone)] |
137 | | pub struct FakeRedisBackend<S: SubscriptionManagerNotify> { |
138 | | /// Contains a list of all of the Redis keys -> fields. |
139 | | pub table: Arc<Mutex<HashMap<String, HashMap<String, Value>>>>, |
140 | | /// TTL (seconds) attached to each key via `EXPIRE`, so tests can |
141 | | /// assert a key was given a bounded lifetime. |
142 | | pub expiries: Arc<Mutex<HashMap<String, i64>>>, |
143 | | subscription_manager: Arc<Mutex<Option<Arc<S>>>>, |
144 | | faults: Arc<Mutex<FaultPolicy>>, |
145 | | } |
146 | | |
147 | | impl<S: SubscriptionManagerNotify + Send + 'static + Sync> Default for FakeRedisBackend<S> { |
148 | 0 | fn default() -> Self { |
149 | 0 | Self::new() |
150 | 0 | } |
151 | | } |
152 | | |
153 | | impl<S: SubscriptionManagerNotify> fmt::Debug for FakeRedisBackend<S> { |
154 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
155 | 0 | f.debug_struct("FakeRedisBackend").finish() |
156 | 0 | } |
157 | | } |
158 | | |
159 | | const FAKE_SCRIPT_SHA: &str = "d89e3573a1f9689c22115e8a41bd332d0c7c2643"; |
160 | | |
161 | | impl<S: SubscriptionManagerNotify + Send + 'static + Sync> FakeRedisBackend<S> { |
162 | 85 | pub fn new() -> Self { |
163 | 85 | Self { |
164 | 85 | table: Arc::new(Mutex::new(HashMap::new())), |
165 | 85 | expiries: Arc::new(Mutex::new(HashMap::new())), |
166 | 85 | subscription_manager: Arc::new(Mutex::new(None)), |
167 | 85 | faults: Arc::new(Mutex::new(FaultPolicy::default())), |
168 | 85 | } |
169 | 85 | } |
170 | | |
171 | | /// Answers the next `times` commands named `command` (`*` for any) |
172 | | /// whose key starts with `key_prefix` with `fault`; `None` is until |
173 | | /// `clear_faults`. |
174 | 14 | pub fn fail_commands( |
175 | 14 | &self, |
176 | 14 | command: &str, |
177 | 14 | key_prefix: Option<&str>, |
178 | 14 | fault: Fault, |
179 | 14 | times: Option<u64>, |
180 | 14 | ) { |
181 | 14 | if times == Some(0) { |
182 | 0 | return; |
183 | 14 | } |
184 | 14 | self.faults.lock().unwrap().rules.push(FaultRule { |
185 | 14 | command: command.to_ascii_uppercase(), |
186 | 14 | key_prefix: key_prefix.map(str::to_string), |
187 | 14 | fault, |
188 | 14 | remaining: times, |
189 | 14 | }); |
190 | 14 | } |
191 | | |
192 | | /// A failover as the client sees it: every connection closes once, |
193 | | /// then writes answer `READONLY` `readonly_writes` times. |
194 | 2 | pub fn failover(&self, readonly_writes: u64) { |
195 | 2 | self.drop_connections(); |
196 | 8 | for command in WRITE_COMMANDS { |
197 | 8 | self.fail_commands(command, None, Fault::ReadOnly, Some(readonly_writes)); |
198 | 8 | } |
199 | 2 | } |
200 | | |
201 | | /// Removes every fault rule, the partition and the latency. |
202 | 1 | pub fn clear_faults(&self) { |
203 | 1 | let mut faults = self.faults.lock().unwrap(); |
204 | 1 | faults.rules.clear(); |
205 | 1 | faults.partitioned = false; |
206 | 1 | faults.latency = None; |
207 | 1 | } |
208 | | |
209 | | /// Closes every open connection and refuses new ones until called |
210 | | /// with `false`. |
211 | 2 | pub fn partition(&self, partitioned: bool) { |
212 | 2 | let mut faults = self.faults.lock().unwrap(); |
213 | 2 | faults.partitioned = partitioned; |
214 | 2 | if partitioned { |
215 | 1 | faults.generation += 1; |
216 | 1 | } |
217 | 2 | } |
218 | | |
219 | | /// Closes every connection made through `port`'s listener and refuses |
220 | | /// new ones there; the data and the other listeners stay. |
221 | 71 | pub fn retire_listener(&self, port: u16) { |
222 | 71 | self.faults.lock().unwrap().retired.insert(port); |
223 | 71 | } |
224 | | |
225 | | /// Closes every open connection once; new ones are accepted. |
226 | 2 | pub fn drop_connections(&self) { |
227 | 2 | self.faults.lock().unwrap().generation += 1; |
228 | 2 | } |
229 | | |
230 | | /// Holds each reply for `latency` (tokio time) before writing it. |
231 | 1 | pub fn set_latency(&self, latency: Option<Duration>) { |
232 | 1 | self.faults.lock().unwrap().latency = latency; |
233 | 1 | } |
234 | | |
235 | | /// Deletes every key starting with `prefix`, as an eviction would. |
236 | | /// Returns how many went. |
237 | 2 | pub fn evict_prefix(&self, prefix: &str) -> usize { |
238 | 2 | let mut table = self.table.lock().unwrap(); |
239 | 2 | let before = table.len(); |
240 | 13 | table2 .retain2 (|key, _| !key.starts_with(prefix)); |
241 | 2 | before - table.len() |
242 | 2 | } |
243 | | |
244 | | /// Commands received since the last reset, by upper-case name. |
245 | 5.55k | pub fn command_counts(&self) -> HashMap<String, u64> { |
246 | 5.55k | self.faults.lock().unwrap().counts.clone() |
247 | 5.55k | } |
248 | | |
249 | | /// Commands received since the fake started. |
250 | 2.00M | pub fn commands_total(&self) -> u64 { |
251 | 2.00M | self.faults.lock().unwrap().total |
252 | 2.00M | } |
253 | | |
254 | 5.55k | pub fn reset_command_counts(&self) { |
255 | 5.55k | self.faults.lock().unwrap().counts.clear(); |
256 | 5.55k | } |
257 | | |
258 | 89 | pub fn set_subscription_manager(&self, subscription_manager: Arc<S>) { |
259 | 89 | self.subscription_manager |
260 | 89 | .lock() |
261 | 89 | .unwrap() |
262 | 89 | .replace(subscription_manager); |
263 | 89 | } |
264 | | |
265 | 88 | async fn dynamic_fake_redis(self, listener: TcpListener, listener_ready_tx: Sender<()>) { |
266 | 88 | let faults = self.faults.clone(); |
267 | 88 | let port = listener.local_addr().map_or(0, |addr| addr.port()); |
268 | 522k | let inner88 = move |buf: &[u8]| -> String { |
269 | 522k | let mut output = String::new(); |
270 | 522k | let mut buf_index = 0; |
271 | | loop { |
272 | 761k | let frame = match decode(&buf[buf_index..]).unwrap() { |
273 | 761k | Some((frame, amt)) => { |
274 | 761k | buf_index += amt; |
275 | 761k | frame |
276 | | } |
277 | | None => { |
278 | 0 | panic!("No frame!"); |
279 | | } |
280 | | }; |
281 | 761k | let (cmd, args) = { |
282 | 761k | if let OwnedFrame::Array(a) = frame { |
283 | 761k | if let OwnedFrame::BulkString(s) = a.first().unwrap() { |
284 | 761k | let args: Vec<_> = a[1..].to_vec(); |
285 | 761k | (str::from_utf8(s).unwrap().to_string(), args) |
286 | | } else { |
287 | 0 | panic!("Array not starting with cmd: {a:?}"); |
288 | | } |
289 | | } else { |
290 | 0 | panic!("Non array cmd: {frame:?}"); |
291 | | } |
292 | | }; |
293 | | |
294 | 761k | let key = args |
295 | 761k | .first() |
296 | 761k | .and_then(|arg| arg.as_bytes()) |
297 | 761k | .and_then(|bytes| str::from_utf8(bytes).ok()); |
298 | | // EVALSHA's key is its third argument, after the sha and the |
299 | | // key count. |
300 | 761k | let key = if cmd.eq_ignore_ascii_case("EVALSHA") { |
301 | 7.91k | args.get(2) |
302 | 7.91k | .and_then(|arg| arg.as_bytes()) |
303 | 7.91k | .and_then(|bytes| str::from_utf8(bytes).ok()) |
304 | | } else { |
305 | 753k | key |
306 | | }; |
307 | 761k | let fault = self |
308 | 761k | .faults |
309 | 761k | .lock() |
310 | 761k | .unwrap() |
311 | 761k | .on_command(&cmd.to_ascii_uppercase(), key); |
312 | 761k | if let Some((reply35 , false)) = &fault { |
313 | 35 | output.push_str(reply); |
314 | 35 | if buf_index == buf.len() { |
315 | 31 | break; |
316 | 4 | } |
317 | 4 | continue; |
318 | 761k | } |
319 | | |
320 | 761k | let ret: Value = match cmd.as_str() { |
321 | 761k | "HELLO" => Value::Map(110 vec!110 [( |
322 | 110 | Value::SimpleString("server".into()), |
323 | 110 | Value::SimpleString("redis".into()), |
324 | 110 | )]), |
325 | 761k | "CLIENT" => { |
326 | | // We can safely ignore these, as it's just setting the library name/version |
327 | 220 | Value::Int(0) |
328 | | } |
329 | 761k | "SCRIPT" => { |
330 | 105 | assert_eq!(args[0], OwnedFrame::BulkString(b"LOAD".to_vec())); |
331 | | |
332 | 105 | let OwnedFrame::BulkString(ref _script) = args[1] else { |
333 | 0 | panic!("Script should be a bulkstring: {args:?}"); |
334 | | }; |
335 | 105 | Value::SimpleString(FAKE_SCRIPT_SHA.to_string()) |
336 | | } |
337 | | |
338 | 761k | "PSUBSCRIBE" => { |
339 | | // This does nothing at the moment, maybe we need to implement it later. |
340 | 210 | Value::Int(0) |
341 | | } |
342 | | |
343 | 761k | "PUBLISH" => { |
344 | 46.6k | if let Some(subscription_manager) = |
345 | 46.6k | self.subscription_manager.lock().unwrap().as_ref() |
346 | | { |
347 | 46.6k | subscription_manager.notify_for_test( |
348 | 46.6k | str::from_utf8(args[1].as_bytes().expect("Notification not bytes")) |
349 | 46.6k | .expect("Notification not UTF-8") |
350 | 46.6k | .into(), |
351 | 46.6k | ); |
352 | 46.6k | Value::Int(1) |
353 | | } else { |
354 | 2 | Value::Int(0) |
355 | | } |
356 | | } |
357 | | |
358 | 714k | "FT.AGGREGATE" => { |
359 | | // The query is either "*" (match all) or @field:{ value }. |
360 | 16.9k | let OwnedFrame::BulkString(ref raw_query) = args[1] else { |
361 | 0 | panic!("Aggregate query should be a string: {args:?}"); |
362 | | }; |
363 | 16.9k | let query = str::from_utf8(raw_query).unwrap(); |
364 | | // The real ft_aggregate caller now passes an explicit |
365 | | // `TIMEOUT <ms>` clause before `LOAD`. Tolerate both |
366 | | // shapes here so this fake doesn't break older callers |
367 | | // and the LOAD-args check still validates the bit we |
368 | | // actually care about. |
369 | 16.9k | let load_offset = if matches!(args.get(2), Some(OwnedFrame::BulkString(b)) if b == b"TIMEOUT") |
370 | | { |
371 | | // Skip "TIMEOUT" and its millisecond argument. |
372 | 16.9k | 4 |
373 | | } else { |
374 | 0 | 2 |
375 | | }; |
376 | | // Lazy implementation making assumptions. |
377 | 16.9k | assert_eq!( |
378 | 16.9k | args[load_offset..load_offset + 4], |
379 | 16.9k | vec![ |
380 | 16.9k | OwnedFrame::BulkString(b"LOAD".to_vec()), |
381 | 16.9k | OwnedFrame::BulkString(b"2".to_vec()), |
382 | 16.9k | OwnedFrame::BulkString(b"data".to_vec()), |
383 | 16.9k | OwnedFrame::BulkString(b"version".to_vec()) |
384 | | ] |
385 | | ); |
386 | 16.9k | let mut results = vec![Value::Int(0)]; |
387 | | |
388 | | // `SORTBY <n> <@field> <ASC|DESC> ...`: honour the first |
389 | | // key so tests can assert the order the scheduler asks |
390 | | // for. Missing sort fields sort first, like RediSearch. |
391 | 16.9k | let sort_spec = args |
392 | 16.9k | .iter() |
393 | 236k | .position16.9k (|a| matches!(a, OwnedFrame::BulkString(b) if b == b"SORTBY"16.9k )) |
394 | 16.9k | .and_then(|i| { |
395 | 16.9k | let field5.71k = args.get(i + 2)?11.2k .as_bytes5.71k ()?0 ; |
396 | 5.71k | let field = str::from_utf8(field).ok()?0 .trim_start_matches('@').to_string(); |
397 | 5.71k | let desc = matches!(args.get(i + 3), Some(OwnedFrame::BulkString(d)) if d == b"DESC"); |
398 | 5.71k | Some((field, desc)) |
399 | 16.9k | }); |
400 | | // Ties come back in key order rather than the |
401 | | // table's hash order, so a run repeats. |
402 | 16.9k | let sorted = |mut rows: Vec<(String, HashMap<String, Value>)>| { |
403 | 5.23M | rows16.9k .sort_by16.9k (|a, b| a.0.cmp(&b.0)); |
404 | 16.9k | let mut rows: Vec<_> = rows.into_iter().map(|(_, f)| f).collect(); |
405 | 16.9k | if let Some((field5.71k , desc5.71k )) = &sort_spec { |
406 | 12.1M | let key5.71k = |r: &HashMap<String, Value>| match r.get(field) { |
407 | 12.1M | Some(Value::BulkString(b)) => b.clone(), |
408 | 0 | _ => Vec::new(), |
409 | 12.1M | }; |
410 | 5.71k | rows.sort_by_key(key); |
411 | 5.71k | if *desc { |
412 | 5.71k | rows.reverse(); |
413 | 5.71k | }0 |
414 | 11.2k | } |
415 | 16.9k | rows |
416 | 16.9k | }; |
417 | | |
418 | 16.9k | if query == "*" { |
419 | | // Wildcard query - return all records that have both data and version fields. |
420 | | // Some entries (e.g., from HSET) may not have version field. |
421 | 0 | let rows: Vec<_> = self |
422 | 0 | .table |
423 | 0 | .lock() |
424 | 0 | .unwrap() |
425 | 0 | .iter() |
426 | 0 | .filter(|(_, f)| { |
427 | 0 | f.contains_key("data") && f.contains_key("version") |
428 | 0 | }) |
429 | 0 | .map(|(k, f)| (k.clone(), f.clone())) |
430 | 0 | .collect(); |
431 | 0 | for fields in sorted(rows) { |
432 | 0 | results.push(Value::Array(vec![ |
433 | 0 | Value::BulkString(b"data".to_vec()), |
434 | 0 | fields["data"].clone(), |
435 | 0 | Value::BulkString(b"version".to_vec()), |
436 | 0 | fields["version"].clone(), |
437 | 0 | ])); |
438 | 0 | } |
439 | | } else { |
440 | | // Field-specific query: @field:{ value } |
441 | 16.9k | assert_eq!(&query[..1], "@"); |
442 | 16.9k | let mut parts = query[1..].split(':'); |
443 | 16.9k | let field = parts.next().expect("No field name"); |
444 | 16.9k | let value = parts.next().expect("No value"); |
445 | 16.9k | let value = value |
446 | 16.9k | .strip_prefix("{ ") |
447 | 16.9k | .and_then(|s| s.strip_suffix(" }")) |
448 | 16.9k | .unwrap_or(value); |
449 | 16.9k | let rows: Vec<_> = self |
450 | 16.9k | .table |
451 | 16.9k | .lock() |
452 | 16.9k | .unwrap() |
453 | 16.9k | .iter() |
454 | 15.8M | .filter16.9k (|(_, f)| { |
455 | 15.8M | f.get(field) |
456 | 15.8M | == Some(&Value::BulkString(value.as_bytes().to_vec())) |
457 | 15.8M | }) |
458 | 485k | .map16.9k (|(k, f)| (k.clone(), f.clone())) |
459 | 16.9k | .collect(); |
460 | 485k | for fields in sorted(rows)16.9k { |
461 | | { |
462 | 485k | let mut record = vec![ |
463 | 485k | Value::BulkString(b"data".to_vec()), |
464 | 485k | fields.get("data").expect("No data field").clone(), |
465 | | ]; |
466 | | // LOAD leaves out fields the hash does not |
467 | | // have, so a non-versioned record comes |
468 | | // back without a version. |
469 | 485k | if let Some(version483k ) = fields.get("version") { |
470 | 483k | record.push(Value::BulkString(b"version".to_vec())); |
471 | 483k | record.push(version.clone()); |
472 | 483k | }1.53k |
473 | 485k | results.push(Value::Array(record)); |
474 | | } |
475 | | } |
476 | | } |
477 | | |
478 | 16.9k | results[0] = |
479 | 16.9k | Value::Int(i64::try_from(results.len() - 1).unwrap_or(i64::MAX)); |
480 | 16.9k | Value::Array(vec![ |
481 | 16.9k | Value::Array(results), |
482 | 16.9k | Value::Int(0), // Means no more items in cursor. |
483 | 16.9k | ]) |
484 | | } |
485 | | |
486 | 697k | "EVALSHA" => { |
487 | 7.89k | assert_eq!( |
488 | 7.89k | args[0], |
489 | 7.89k | OwnedFrame::BulkString(FAKE_SCRIPT_SHA.as_bytes().to_vec()) |
490 | | ); |
491 | 7.89k | assert_eq!(args[1], OwnedFrame::BulkString(b"1".to_vec())); |
492 | 7.89k | let mut value: HashMap<_, Value> = HashMap::new(); |
493 | 7.89k | value.insert( |
494 | 7.89k | "data".into(), |
495 | 7.89k | Value::BulkString(args[5].as_bytes().unwrap().to_vec()), |
496 | | ); |
497 | 23.4k | for pair in args[6..]7.89k .chunks7.89k (2) { |
498 | 23.4k | value.insert( |
499 | 23.4k | str::from_utf8(pair[0].as_bytes().expect("Field name not bytes")) |
500 | 23.4k | .expect("Unable to parse field name as string") |
501 | 23.4k | .into(), |
502 | 23.4k | Value::BulkString(pair[1].as_bytes().unwrap().to_vec()), |
503 | 23.4k | ); |
504 | 23.4k | } |
505 | 7.89k | let mut ret: Option<Value> = None; |
506 | 7.89k | let key: String = |
507 | 7.89k | str::from_utf8(args[2].as_bytes().expect("Key not bytes")) |
508 | 7.89k | .expect("Key cannot be parsed as string") |
509 | 7.89k | .into(); |
510 | 7.89k | let expected_existing_version: i64 = |
511 | 7.89k | str::from_utf8(args[3].as_bytes().unwrap()) |
512 | 7.89k | .unwrap() |
513 | 7.89k | .parse() |
514 | 7.89k | .expect("Unable to parse existing version field"); |
515 | 7.89k | let expiry: i64 = str::from_utf8(args[4].as_bytes().unwrap()) |
516 | 7.89k | .unwrap() |
517 | 7.89k | .parse() |
518 | 7.89k | .expect("Unable to parse expiry field"); |
519 | 7.89k | trace!( |
520 | | key, |
521 | | expected_existing_version, |
522 | | expiry, |
523 | | ?value, |
524 | | "Want to insert with EVALSHA" |
525 | | ); |
526 | 7.89k | let version = match self.table.lock().unwrap().entry(key.clone()) { |
527 | 2.87k | Entry::Occupied(mut occupied_entry) => { |
528 | 2.87k | let version = occupied_entry |
529 | 2.87k | .get() |
530 | 2.87k | .get("version") |
531 | 2.87k | .expect("No version field"); |
532 | 2.87k | let Value::BulkString(version_bytes) = version else { |
533 | 0 | panic!("Non-bulkstring version: {version:?}"); |
534 | | }; |
535 | 2.87k | let version_int: i64 = str::from_utf8(version_bytes) |
536 | 2.87k | .expect("Version field not valid string") |
537 | 2.87k | .parse() |
538 | 2.87k | .expect("Unable to parse version field"); |
539 | 2.87k | if version_int == expected_existing_version { |
540 | 2.86k | let new_version = version_int + 1; |
541 | 2.86k | debug!(%key, %new_version, "Version update"); |
542 | 2.86k | value.insert( |
543 | 2.86k | "version".into(), |
544 | 2.86k | Value::BulkString( |
545 | 2.86k | format!("{new_version}").as_bytes().to_vec(), |
546 | 2.86k | ), |
547 | | ); |
548 | 2.86k | occupied_entry.insert(value); |
549 | 2.86k | new_version |
550 | | } else { |
551 | | // Version mismatch. |
552 | 9 | debug!(%key, %version_int, %expected_existing_version, "Version mismatch"); |
553 | 9 | ret = Some(Value::Array(vec![ |
554 | 9 | Value::Int(0), |
555 | 9 | Value::Int(version_int), |
556 | 9 | ])); |
557 | 9 | -1 |
558 | | } |
559 | | } |
560 | 5.02k | Entry::Vacant(vacant_entry) => { |
561 | 5.02k | if expected_existing_version != 0 { |
562 | | // Version mismatch. |
563 | 1 | debug!(%key, %expected_existing_version, "Version mismatch, expected zero"); |
564 | 1 | ret = Some(Value::Array(vec![Value::Int(0), Value::Int(0)])); |
565 | 1 | -1 |
566 | | } else { |
567 | 5.02k | debug!(%key, "Version insert"); |
568 | 5.02k | value |
569 | 5.02k | .insert("version".into(), Value::BulkString(b"1".to_vec())); |
570 | 5.02k | vacant_entry.insert_entry(value); |
571 | 5.02k | 1 |
572 | | } |
573 | | } |
574 | | }; |
575 | 7.89k | if let Some(r10 ) = ret { |
576 | 10 | r |
577 | | } else { |
578 | 7.88k | Value::Array(vec![Value::Int(1), Value::Int(version)]) |
579 | | } |
580 | | } |
581 | | |
582 | 689k | "HMSET" => { |
583 | 38.8k | let mut values = HashMap::new(); |
584 | 38.8k | assert_eq!( |
585 | 38.8k | (args.len() - 1).rem_euclid(2), |
586 | | 0, |
587 | | "Non-even args for hmset: {args:?}" |
588 | | ); |
589 | 38.8k | let (chunks, rest) = args[1..].as_chunks::<2>(); |
590 | 38.8k | assert_eq!(rest, [], "Uneven hmset args"); |
591 | 40.1k | for [key, value] in chunks38.8k { |
592 | 40.1k | let key_name: String = |
593 | 40.1k | str::from_utf8(key.as_bytes().expect("Key argument is not bytes")) |
594 | 40.1k | .expect("Unable to parse key as string") |
595 | 40.1k | .into(); |
596 | 40.1k | values.insert( |
597 | 40.1k | key_name, |
598 | 40.1k | Value::BulkString(value.as_bytes().unwrap().to_vec()), |
599 | 40.1k | ); |
600 | 40.1k | } |
601 | 38.8k | let key = |
602 | 38.8k | str::from_utf8(args[0].as_bytes().expect("Key argument is not bytes")) |
603 | 38.8k | .expect("Unable to parse key as string") |
604 | 38.8k | .into(); |
605 | 38.8k | debug!(%key, ?values, "Inserting with HMSET"); |
606 | 38.8k | self.table.lock().unwrap().insert(key, values); |
607 | 38.8k | Value::Okay |
608 | | } |
609 | | |
610 | 650k | "EXPIRE" => { |
611 | | // `EXPIRE key seconds`: record the TTL; return 1 |
612 | | // if the key existed, else 0 (as real Redis does). |
613 | 38.7k | let key_name = |
614 | 38.7k | str::from_utf8(args[0].as_bytes().expect("Key argument is not bytes")) |
615 | 38.7k | .expect("Unable to parse key name") |
616 | 38.7k | .to_string(); |
617 | 38.7k | let seconds = str::from_utf8( |
618 | 38.7k | args[1].as_bytes().expect("EXPIRE seconds is not bytes"), |
619 | 38.7k | ) |
620 | 38.7k | .expect("EXPIRE seconds is not utf8") |
621 | 38.7k | .parse::<i64>() |
622 | 38.7k | .expect("EXPIRE seconds is not an integer"); |
623 | 38.7k | let exists = self.table.lock().unwrap().contains_key(&key_name); |
624 | 38.7k | if exists { |
625 | 38.7k | self.expiries.lock().unwrap().insert(key_name, seconds); |
626 | 38.7k | Value::Int(1) |
627 | | } else { |
628 | 0 | Value::Int(0) |
629 | | } |
630 | | } |
631 | | |
632 | 612k | "HMGET" => { |
633 | 612k | let key_name = |
634 | 612k | str::from_utf8(args[0].as_bytes().expect("Key argument is not bytes")) |
635 | 612k | .expect("Unable to parse key name"); |
636 | | |
637 | 612k | if let Some(fields108k ) = self.table.lock().unwrap().get(key_name) { |
638 | 108k | trace!(%key_name, keys = ?fields718 .keys718 (), "Getting keys with HMGET, some keys"); |
639 | 108k | let mut result = vec![]; |
640 | 217k | for key in &args[1..]108k { |
641 | 217k | let field_name = str::from_utf8( |
642 | 217k | key.as_bytes().expect("Field argument is not bytes"), |
643 | | ) |
644 | 217k | .expect("Unable to parse requested field"); |
645 | 217k | if let Some(value179k ) = fields.get(field_name) { |
646 | 179k | result.push(value.clone()); |
647 | 179k | } else { |
648 | 38.1k | debug!(%key_name, %field_name, "Missing field"); |
649 | 38.1k | result.push(Value::Nil); |
650 | | } |
651 | | } |
652 | 108k | Value::Array(result) |
653 | | } else { |
654 | 503k | trace!(%key_name, "Getting keys with HMGET, empty"); |
655 | 503k | let null_count = i64::try_from(args.len() - 1).unwrap(); |
656 | 503k | Value::Array(vec![Value::Nil, Value::Int(null_count)]) |
657 | | } |
658 | | } |
659 | 0 | actual => { |
660 | 0 | panic!("Mock command not implemented! {actual:?}"); |
661 | | } |
662 | | }; |
663 | | |
664 | 761k | match fault { |
665 | 1 | Some((reply, true)) => output.push_str(&reply), |
666 | 761k | _ => arg_as_string(&mut output, ret), |
667 | | } |
668 | 761k | if buf_index == buf.len() { |
669 | 522k | break; |
670 | 238k | } |
671 | | } |
672 | 522k | output |
673 | 522k | }; |
674 | 88 | listener_ready_tx |
675 | 88 | .send(()) |
676 | 88 | .expect("Expected successful send"); |
677 | | loop { |
678 | 1.11k | let Ok((mut stream1.02k , _)) = listener.accept().await else { |
679 | 0 | panic!("accept error"); |
680 | | }; |
681 | 110 | let generation = { |
682 | 1.02k | let faults = faults.lock().unwrap(); |
683 | 1.02k | if faults.partitioned || faults.retired1.02k .contains1.02k (&port1.02k ) { |
684 | 914 | drop(stream); |
685 | 914 | continue; |
686 | 110 | } |
687 | 110 | faults.generation |
688 | | }; |
689 | 110 | let handler = inner.clone(); |
690 | 110 | let faults = faults.clone(); |
691 | 110 | background_spawn!("dynamic_fake_redis_connection", async move { |
692 | 110 | let mut buf = vec![0; 8192]; |
693 | | // Bytes read but not yet a whole command: a command can span |
694 | | // two reads. |
695 | 110 | let mut pending: Vec<u8> = Vec::new(); |
696 | | loop { |
697 | 523k | let Ok(read522k ) = stream.read(&mut buf).await else { |
698 | 0 | return; |
699 | | }; |
700 | | { |
701 | 522k | let faults = faults.lock().unwrap(); |
702 | 522k | if read == 0 |
703 | 522k | || faults.generation != generation |
704 | 522k | || faults.retired.contains(&port) |
705 | | { |
706 | 22 | return; |
707 | 522k | } |
708 | | } |
709 | 522k | pending.extend_from_slice(&buf[..read]); |
710 | 522k | let complete = complete_frames_len(&pending); |
711 | 522k | if complete == 0 { |
712 | 0 | continue; |
713 | 522k | } |
714 | 522k | let commands: Vec<u8> = pending.drain(..complete).collect(); |
715 | 522k | let output = handler(&commands); |
716 | 522k | let latency = faults.lock().unwrap().latency; |
717 | 522k | if let Some(latency5 ) = latency { |
718 | 5 | tokio::time::sleep(latency).await; |
719 | 522k | } |
720 | 522k | if faults.lock().unwrap().generation != generation { |
721 | 0 | return; |
722 | 522k | } |
723 | 522k | if !output.is_empty() && stream.write_all(output.as_bytes()).await.is_err() { |
724 | 0 | return; |
725 | 522k | } |
726 | | } |
727 | 22 | }); |
728 | | } |
729 | | } |
730 | | |
731 | 88 | pub async fn run(self) -> u16 { |
732 | 88 | let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); |
733 | 88 | let port = listener.local_addr().unwrap().port(); |
734 | 88 | info!("Using port {port}"); |
735 | | |
736 | 88 | let (listener_ready_tx, listener_ready_rx) = oneshot::channel::<()>(); |
737 | | |
738 | 88 | background_spawn!("listener", async move { |
739 | 88 | self.dynamic_fake_redis(listener, listener_ready_tx).await; |
740 | 0 | }); |
741 | | |
742 | 88 | listener_ready_rx |
743 | 88 | .await |
744 | 88 | .expect("Expected successful listener boot"); |
745 | | |
746 | 88 | port |
747 | 88 | } |
748 | | } |