Coverage Report

Created: 2026-09-18 20:40

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/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
1.01k
pub(crate) fn arg_as_string(output: &mut String, arg: Value) {
33
1.01k
    match arg {
34
120
        Value::SimpleString(s) => {
35
120
            write!(output, "+{s}\r\n").unwrap();
36
120
        }
37
147
        Value::Okay => {
38
147
            write!(output, "+OK\r\n").unwrap();
39
147
        }
40
214
        Value::BulkString(s) => {
41
214
            write!(
42
214
                output,
43
214
                "${}\r\n{}\r\n",
44
214
                s.len(),
45
214
                str::from_utf8(&s).unwrap()
46
214
            )
47
214
            .unwrap();
48
214
        }
49
216
        Value::Int(v) => {
50
216
            write!(output, ":{v}\r\n").unwrap();
51
216
        }
52
196
        Value::Array(values) => {
53
196
            write!(output, "*{}\r\n", values.len()).unwrap();
54
384
            for value in 
values196
{
55
384
                arg_as_string(output, value);
56
384
            }
57
        }
58
66
        Value::Map(values) => {
59
66
            write!(output, "%{}\r\n", values.len()).unwrap();
60
87
            for (key, value) in 
values66
{
61
87
                arg_as_string(output, key);
62
87
                arg_as_string(output, value);
63
87
            }
64
        }
65
53
        Value::Nil => {
66
53
            write!(output, "_\r\n").unwrap();
67
53
        }
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
1.01k
}
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
29
pub(crate) async fn fake_redis_internal<H>(
203
29
    listener: TcpListener,
204
29
    listener_ready_tx: Sender<()>,
205
29
    handlers: Vec<H>,
206
29
) where
207
29
    H: Fn(&[u8]) -> String + Send + Clone + 'static + Sync,
208
29
{
209
29
    let mut handler_iter = handlers.iter().cloned().cycle();
210
29
    info!(
211
        "Waiting for connection on {}",
212
29
        listener.local_addr().unwrap()
213
    );
214
29
    listener_ready_tx
215
29
        .send(())
216
29
        .expect("Expected successful send");
217
    loop {
218
72
        let Ok((
mut stream43
, _)) = listener.accept().await else {
219
0
            error!("accept error");
220
0
            panic!("error");
221
        };
222
43
        info!("Accepted new connection");
223
43
        let local_handler = handler_iter.next().unwrap();
224
43
        background_spawn!("thread", async move {
225
            loop {
226
99.7k
                let mut buf = vec![0; 8192];
227
99.7k
                let 
res99.7k
= stream.read(&mut buf).await.
unwrap99.7k
();
228
99.7k
                if res != 0 {
229
533
                    let output = local_handler(&buf[..res]);
230
533
                    if !output.is_empty() {
231
271
                        stream.write_all(output.as_bytes()).await.unwrap();
232
262
                    }
233
99.2k
                }
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
327
            move |buf: &[u8]| -> String {
252
327
                let str_buf = String::from_utf8_lossy(buf).into_owned();
253
1.22k
                for (key, value) in 
&values327
{
254
1.22k
                    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.15k
                    }
258
                }
259
262
                warn!(
260
                    "Unknown command: {}",
261
262
                    str_buf.chars().take(1000).collect::<String>()
262
                );
263
262
                String::new()
264
327
            }
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
}