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

Source
// Copyright 2026 Trace Machina, Inc. 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 = "d89e3573a1f9689c22115e8a41bd332d0c7c2643";

#[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<()>)0 {
        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;
            // The replies queued since a MULTI, answered by its EXEC.
            let mut queued: Option<Vec<Value>> = None;
            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,
                    "TTL" => Either::Left(Value::Int(60))2,
                    "PERSIST" => Either::Left(Value::Int(1))1,
                    "SET" | "RENAME"8 | "HMSET"7 => {
                        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)])),
                    "MULTI" => {
                        queued = Some(Vec::new());
                        Either::Left(Value::Okay)
                    }
                    "EXEC" => Either::Left(Value::Array(1queued1.take().unwrap_or_default())),
                    "EXPIRE" => {
                        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:?}");
                    }
                };

                let ret = match (ret, queued.as_mut(), cmd.as_str()) {
                    (Either::Left(v2), Some(queued2), cmd2) if cmd != "MULTI" && cmd2 != "EXEC"2 => {
                        queued.push(v);
                        Either::Left(Value::SimpleString("QUEUED".into()))
                    }
                    (ret, _, _) => ret,
                };
                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) -> u16 {
        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
    }
}

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::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 = "d89e3573a1f9689c22115e8a41bd332d0c7c2643";
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<()>)0 {
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	29	let inner3 = move \|buf: &[u8]\| -> String {
56	29	let mut output = String::new();
57	29	let mut buf_index = 0;
58		// The replies queued since a MULTI, answered by its EXEC.
59	29	let mut queued: Option<Vec<Value>> = None;
60		loop {
61	55	let frame = match decode(&buf[buf_index..]).unwrap() {
62	55	Some((frame, amt)) => {
63	55	buf_index += amt;
64	55	frame
65		}
66		None => {
67	0	panic!("No frame!");
68		}
69		};
70	55	let (cmd, args) = {
71	55	if let OwnedFrame::Array(a) = frame {
72	55	if let OwnedFrame::BulkString(s) = a.first().unwrap() {
73	55	let args: Vec<_> = a[1..].to_vec();
74	55	(str::from_utf8(s).unwrap().to_string(), args)
75		} else {
76	0	panic!("Array not starting with cmd: {a:?}");
77		}
78		} else {
79	0	panic!("Non array cmd: {frame:?}");
80		}
81		};
82
83	55	let ret: Either<Value, String> = match cmd.as_str() {
84	55	"HELLO" => Either::Left(Value::Map(10 vec!10 [(
85	10	Value::SimpleString("server".into()),
86	10	Value::SimpleString("redis".into()),
87	10	)])),
88	45	"CLIENT" => {
89		// We can safely ignore these, as it's just setting the library name/version
90	20	Either::Left(Value::Int(0))
91		}
92	25	"SCRIPT" => {
93	5	assert_eq!(args[0], OwnedFrame::BulkString(b"LOAD".to_vec()));
94
95	5	let OwnedFrame::BulkString(ref _script) = args[1] else {
96	0	panic!("Script should be a bulkstring: {args:?}");
97		};
98	5	Either::Left(Value::SimpleString(FAKE_SCRIPT_SHA.to_string()))
99		}
100	20	"ROLE" => Either::Left(Value::Array(5 vec!5 [
101	5	Value::BulkString(b"master".to_vec()),
102	5	Value::Int(0),
103	5	Value::Array(vec![]),
104	5	])),
105	15	"SETRANGE" => {
106	1	let value = self.readonly_triggered.load(Ordering::Relaxed);
107	1	if value {
108	1	Either::Left(Value::Int(5))
109		} else {
110	0	self.readonly_triggered.store(true, Ordering::Relaxed);
111	0	Either::Right(readonly_err.clone())
112		}
113		}
114	14	"STRLEN" => Either::Left(Value::Int(5))1 ,
115	13	"TTL" => Either::Left(Value::Int(60))2 ,
116	11	"PERSIST" => Either::Left(Value::Int(1))1 ,
117	10	"SET" \| "RENAME"8 \| "HMSET"7 => {
118	5	let value = self.readonly_triggered.load(Ordering::Relaxed);
119	5	if value {
120	3	Either::Left(Value::Okay)
121		} else {
122	2	self.readonly_triggered.store(true, Ordering::Relaxed);
123	2	Either::Right(readonly_err.clone())
124		}
125		}
126	5	"EVALSHA" => Either::Left(Value::Array(1 vec!1 [Value::Int(1), Value::Int(0)])),
127	4	"MULTI" => {
128	1	queued = Some(Vec::new());
129	1	Either::Left(Value::Okay)
130		}
131	3	"EXEC" => Either::Left(Value::Array(1 queued1 .take().unwrap_or_default())),
132	2	"EXPIRE" => {
133	2	let value = self.readonly_triggered.load(Ordering::Relaxed);
134	2	if value {
135	2	Either::Left(Value::Int(1))
136		} else {
137	0	self.readonly_triggered.store(true, Ordering::Relaxed);
138	0	Either::Right(readonly_err.clone())
139		}
140		}
141	0	actual => {
142	0	panic!("Mock command not implemented! {actual:?}");
143		}
144		};
145
146	55	let ret = match (ret, queued.as_mut(), cmd.as_str()) {
147	3	(Either::Left( v2 ), Some( queued2 ), cmd2 ) if cmd != "MULTI" && cmd2 != "EXEC "2 => {
148	2	queued.push(v);
149	2	Either::Left(Value::SimpleString("QUEUED".into()))
150		}
151	53	(ret, _, _) => ret,
152		};
153	55	match ret {
154	53	Either::Left(v) => {
155	53	arg_as_string(&mut output, v);
156	53	}
157	2	Either::Right(s) => {
158	2	write!(&mut output, "{s}").unwrap();
159	2	}
160		}
161
162	55	if buf_index == buf.len() {
163	29	break;
164	26	}
165		}
166	29	output
167	29	};
168	3	fake_redis_internal(listener, listener_ready_tx, vec![inner]).await;
169	0	}
170
171	3	pub async fn run(self) -> u16 {
172	3	let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
173	3	let port = listener.local_addr().unwrap().port();
174	3	info!("Using port {port}");
175
176	3	let (listener_ready_tx, listener_ready_rx) = oneshot::channel::<()>();
177
178	3	background_spawn!("listener", async move {
179	3	self.dynamic_fake_redis(listener, listener_ready_tx).await;
180	0	});
181
182	3	listener_ready_rx
183	3	.await
184	3	.expect("Expected successful listener boot");
185
186	3	port
187	3	}
188		}

Coverage Report

Created: 2026-11-13 23:29