/build/source/nativelink-redis-tester/src/read_only_redis.rs

Source
// Copyright 2026 The NativeLink Authors. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use core::fmt::Write;
use core::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use either::Either;
use nativelink_util::background_spawn;
use redis::Value;
use redis_protocol::resp2::decode::decode;
use redis_protocol::resp2::types::OwnedFrame;
use tokio::net::TcpListener;
use tokio::sync::oneshot::{self, Sender};
use tracing::info;

use crate::fake_redis::{arg_as_string, fake_redis_internal};

const FAKE_SCRIPT_SHA: &str = "5148c724ce419ea27d1971dcb61c111dbbc6b63e";

#[derive(Clone, Debug)]
pub struct ReadOnlyRedis {
    // The first time we hit SETRANGE/HMSET, we output a ReadOnly. Next time, we assume we're reconnected and do correct values
    readonly_triggered: Arc<AtomicBool>,
}

impl Default for ReadOnlyRedis {
    fn default() -> Self {
        Self::new()
    }
}

impl ReadOnlyRedis {
    pub fn new() -> Self {
        Self {
            readonly_triggered: Arc::new(AtomicBool::new(false)),
        }
    }

    async fn dynamic_fake_redis(self, listener: TcpListener, listener_ready_tx: Sender<()>) {
        let readonly_err_str = "READONLY You can't write against a read only replica.";
        let readonly_err = format!("!{}\r\n{readonly_err_str}\r\n", readonly_err_str.len());

        let inner3 = move |buf: &[u8]| -> String {
            let mut output = String::new();
            let mut buf_index = 0;
            loop {
                let frame = match decode(&buf[buf_index..]).unwrap() {
                    Some((frame, amt)) => {
                        buf_index += amt;
                        frame
                    }
                    None => {
                        panic!("No frame!");
                    }
                };
                let (cmd, args) = {
                    if let OwnedFrame::Array(a) = frame {
                        if let OwnedFrame::BulkString(s) = a.first().unwrap() {
                            let args: Vec<_> = a[1..].to_vec();
                            (str::from_utf8(s).unwrap().to_string(), args)
                        } else {
                            panic!("Array not starting with cmd: {a:?}");
                        }
                    } else {
                        panic!("Non array cmd: {frame:?}");
                    }
                };

                let ret: Either<Value, String> = match cmd.as_str() {
                    "HELLO" => Either::Left(Value::Map(10vec!10[(
                        Value::SimpleString("server".into()),
                        Value::SimpleString("redis".into()),
                    )])),
                    "CLIENT" => {
                        // We can safely ignore these, as it's just setting the library name/version
                        Either::Left(Value::Int(0))
                    }
                    "SCRIPT" => {
                        assert_eq!(args[0], OwnedFrame::BulkString(b"LOAD".to_vec()));

                        let OwnedFrame::BulkString(ref _script) = args[1] else {
                            panic!("Script should be a bulkstring: {args:?}");
                        };
                        Either::Left(Value::SimpleString(FAKE_SCRIPT_SHA.to_string()))
                    }
                    "ROLE" => Either::Left(Value::Array(5vec!5[
                        Value::BulkString(b"master".to_vec()),
                        Value::Int(0),
                        Value::Array(vec![]),
                    ])),
                    "SETRANGE" => {
                        let value = self.readonly_triggered.load(Ordering::Relaxed);
                        if value {
                            Either::Left(Value::Int(5))
                        } else {
                            self.readonly_triggered.store(true, Ordering::Relaxed);
                            Either::Right(readonly_err.clone())
                        }
                    }
                    "STRLEN" => Either::Left(Value::Int(5))1,
                    "RENAME" | "HMSET"4 => {
                        let value = self.readonly_triggered.load(Ordering::Relaxed);
                        if value {
                            Either::Left(Value::Okay)
                        } else {
                            self.readonly_triggered.store(true, Ordering::Relaxed);
                            Either::Right(readonly_err.clone())
                        }
                    }
                    "EVALSHA" => Either::Left(Value::Array(1vec!1[Value::Int(1), Value::Int(0)])),
                    "EXPIRE" => {
                        assert_eq!(args[1], OwnedFrame::BulkString(b"60".to_vec()));
                        let value = self.readonly_triggered.load(Ordering::Relaxed);
                        if value {
                            Either::Left(Value::Int(1))
                        } else {
                            self.readonly_triggered.store(true, Ordering::Relaxed);
                            Either::Right(readonly_err.clone())
                        }
                    }
                    actual => {
                        panic!("Mock command not implemented! {actual:?}");
                    }
                };

                match ret {
                    Either::Left(v) => {
                        arg_as_string(&mut output, v);
                    }
                    Either::Right(s) => {
                        write!(&mut output, "{s}").unwrap();
                    }
                }

                if buf_index == buf.len() {
                    break;
                }
            }
            output
        };
        fake_redis_internal(listener, listener_ready_tx, vec![inner]).await;
    }

    pub async fn run(self) -> u160 {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();
        info!("Using port {port}");

        let (listener_ready_tx, listener_ready_rx) = oneshot::channel::<()>();

        background_spawn!("listener", async move {
            self.dynamic_fake_redis(listener, listener_ready_tx).await;
        });

        listener_ready_rx
            .await
            .expect("Expected successful listener boot");

        port
    }
}

Coverage Report

Created: 2026-07-21 15:28

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::sync::atomic::{AtomicBool, Ordering};
17		use std::sync::Arc;
18
19		use either::Either;
20		use nativelink_util::background_spawn;
21		use redis::Value;
22		use redis_protocol::resp2::decode::decode;
23		use redis_protocol::resp2::types::OwnedFrame;
24		use tokio::net::TcpListener;
25		use tokio::sync::oneshot::{self, Sender};
26		use tracing::info;
27
28		use crate::fake_redis::{arg_as_string, fake_redis_internal};
29
30		const FAKE_SCRIPT_SHA: &str = "5148c724ce419ea27d1971dcb61c111dbbc6b63e";
31
32		#[derive(Clone, Debug)]
33		pub struct ReadOnlyRedis {
34		// The first time we hit SETRANGE/HMSET, we output a ReadOnly. Next time, we assume we're reconnected and do correct values
35		readonly_triggered: Arc<AtomicBool>,
36		}
37
38		impl Default for ReadOnlyRedis {
39	0	fn default() -> Self {
40	0	Self::new()
41	0	}
42		}
43
44		impl ReadOnlyRedis {
45	3	pub fn new() -> Self {
46	3	Self {
47	3	readonly_triggered: Arc::new(AtomicBool::new(false)),
48	3	}
49	3	}
50
51	3	async fn dynamic_fake_redis(self, listener: TcpListener, listener_ready_tx: Sender<()>) {
52	3	let readonly_err_str = "READONLY You can't write against a read only replica.";
53	3	let readonly_err = format!("!{}\r\n{readonly_err_str}\r\n", readonly_err_str.len());
54
55	28	let inner3 = move \|buf: &[u8]\| -> String {
56	28	let mut output = String::new();
57	28	let mut buf_index = 0;
58		loop {
59	48	let frame = match decode(&buf[buf_index..]).unwrap() {
60	48	Some((frame, amt)) => {
61	48	buf_index += amt;
62	48	frame
63		}
64		None => {
65	0	panic!("No frame!");
66		}
67		};
68	48	let (cmd, args) = {
69	48	if let OwnedFrame::Array(a) = frame {
70	48	if let OwnedFrame::BulkString(s) = a.first().unwrap() {
71	48	let args: Vec<_> = a[1..].to_vec();
72	48	(str::from_utf8(s).unwrap().to_string(), args)
73		} else {
74	0	panic!("Array not starting with cmd: {a:?}");
75		}
76		} else {
77	0	panic!("Non array cmd: {frame:?}");
78		}
79		};
80
81	48	let ret: Either<Value, String> = match cmd.as_str() {
82	48	"HELLO" => Either::Left(Value::Map(10 vec!10 [(
83	10	Value::SimpleString("server".into()),
84	10	Value::SimpleString("redis".into()),
85	10	)])),
86	38	"CLIENT" => {
87		// We can safely ignore these, as it's just setting the library name/version
88	20	Either::Left(Value::Int(0))
89		}
90	18	"SCRIPT" => {
91	5	assert_eq!(args[0], OwnedFrame::BulkString(b"LOAD".to_vec()));
92
93	5	let OwnedFrame::BulkString(ref _script) = args[1] else {
94	0	panic!("Script should be a bulkstring: {args:?}");
95		};
96	5	Either::Left(Value::SimpleString(FAKE_SCRIPT_SHA.to_string()))
97		}
98	13	"ROLE" => Either::Left(Value::Array(5 vec!5 [
99	5	Value::BulkString(b"master".to_vec()),
100	5	Value::Int(0),
101	5	Value::Array(vec![]),
102	5	])),
103	8	"SETRANGE" => {
104	2	let value = self.readonly_triggered.load(Ordering::Relaxed);
105	2	if value {
106	1	Either::Left(Value::Int(5))
107		} else {
108	1	self.readonly_triggered.store(true, Ordering::Relaxed);
109	1	Either::Right(readonly_err.clone())
110		}
111		}
112	6	"STRLEN" => Either::Left(Value::Int(5))1 ,
113	5	"RENAME" \| "HMSET"4 => {
114	3	let value = self.readonly_triggered.load(Ordering::Relaxed);
115	3	if value {
116	2	Either::Left(Value::Okay)
117		} else {
118	1	self.readonly_triggered.store(true, Ordering::Relaxed);
119	1	Either::Right(readonly_err.clone())
120		}
121		}
122	2	"EVALSHA" => Either::Left(Value::Array(1 vec!1 [Value::Int(1), Value::Int(0)])),
123	1	"EXPIRE" => {
124	1	assert_eq!(args[1], OwnedFrame::BulkString(b"60".to_vec()));
125	1	let value = self.readonly_triggered.load(Ordering::Relaxed);
126	1	if value {
127	1	Either::Left(Value::Int(1))
128		} else {
129	0	self.readonly_triggered.store(true, Ordering::Relaxed);
130	0	Either::Right(readonly_err.clone())
131		}
132		}
133	0	actual => {
134	0	panic!("Mock command not implemented! {actual:?}");
135		}
136		};
137
138	48	match ret {
139	46	Either::Left(v) => {
140	46	arg_as_string(&mut output, v);
141	46	}
142	2	Either::Right(s) => {
143	2	write!(&mut output, "{s}").unwrap();
144	2	}
145		}
146
147	48	if buf_index == buf.len() {
148	28	break;
149	20	}
150		}
151	28	output
152	28	};
153	3	fake_redis_internal(listener, listener_ready_tx, vec![inner]).await;
154	0	}
155
156	3	pub async fn run(self) -> u160 {
157	3	let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
158	3	let port = listener.local_addr().unwrap().port();
159	3	info!("Using port {port}");
160
161	3	let (listener_ready_tx, listener_ready_rx) = oneshot::channel::<()>();
162
163	3	background_spawn!("listener", async move {
164	3	self.dynamic_fake_redis(listener, listener_ready_tx).await;
165	0	});
166
167	3	listener_ready_rx
168	3	.await
169	3	.expect("Expected successful listener boot");
170
171	3	port
172	3	}
173		}