/build/source/src/bin/redis_store_tester.rs
Line | Count | Source |
1 | | use core::pin::Pin; |
2 | | use core::sync::atomic::{AtomicUsize, Ordering}; |
3 | | use core::time::Duration; |
4 | | use std::borrow::Cow; |
5 | | use std::env; |
6 | | use std::sync::{Arc, RwLock}; |
7 | | |
8 | | use bytes::Bytes; |
9 | | use clap::{Parser, ValueEnum}; |
10 | | use futures::TryStreamExt; |
11 | | use nativelink_config::stores::{RedisMode, RedisSpec}; |
12 | | use nativelink_error::{Code, Error, ResultExt}; |
13 | | use nativelink_store::redis_store::RedisStore; |
14 | | use nativelink_util::buf_channel::make_buf_channel_pair; |
15 | | use nativelink_util::store_trait::{ |
16 | | RemoveItemCallback, SchedulerCurrentVersionProvider, SchedulerIndexProvider, SchedulerStore, |
17 | | SchedulerStoreDataProvider, SchedulerStoreDecodeTo, SchedulerStoreKeyProvider, StoreDriver, |
18 | | StoreKey, StoreLike, TrueValue, UploadSizeInfo, |
19 | | }; |
20 | | use nativelink_util::telemetry::init_tracing; |
21 | | use nativelink_util::{background_spawn, spawn}; |
22 | | use rand::Rng; |
23 | | use tokio::time::sleep; |
24 | | use tracing::{error, info}; |
25 | | |
26 | | // Define test structures that implement the scheduler traits |
27 | | #[derive(Debug, Clone, PartialEq)] |
28 | | struct TestSchedulerData { |
29 | | key: String, |
30 | | content: String, |
31 | | version: i64, |
32 | | } |
33 | | |
34 | | #[derive(Debug)] |
35 | | struct TestSchedulerReturn { |
36 | | version: i64, |
37 | | } |
38 | | |
39 | | impl SchedulerStoreKeyProvider for TestSchedulerData { |
40 | | type Versioned = TrueValue; // Using versioned storage |
41 | | |
42 | 0 | fn get_key(&self) -> StoreKey<'static> { |
43 | 0 | StoreKey::Str(Cow::Owned(self.key.clone())) |
44 | 0 | } |
45 | | } |
46 | | |
47 | | impl SchedulerStoreDataProvider for TestSchedulerData { |
48 | 0 | fn try_into_bytes(self) -> Result<Bytes, Error> { |
49 | 0 | Ok(Bytes::from(self.content)) |
50 | 0 | } |
51 | | |
52 | 0 | fn get_indexes(&self) -> Result<Vec<(&'static str, Bytes)>, Error> { |
53 | | // Add some test indexes - need to use 'static strings |
54 | 0 | Ok(vec![ |
55 | 0 | ("test_index", Bytes::from("test_value")), |
56 | 0 | ( |
57 | 0 | "content_prefix", |
58 | 0 | Bytes::from(self.content.chars().take(10).collect::<String>()), |
59 | 0 | ), |
60 | 0 | ]) |
61 | 0 | } |
62 | | } |
63 | | |
64 | | impl SchedulerStoreDecodeTo for TestSchedulerData { |
65 | | type DecodeOutput = TestSchedulerReturn; |
66 | | |
67 | 0 | fn decode(version: i64, _data: Bytes) -> Result<Self::DecodeOutput, Error> { |
68 | 0 | Ok(TestSchedulerReturn { version }) |
69 | 0 | } |
70 | | } |
71 | | |
72 | | impl SchedulerCurrentVersionProvider for TestSchedulerData { |
73 | 0 | fn current_version(&self) -> i64 { |
74 | 0 | self.version |
75 | 0 | } |
76 | | } |
77 | | |
78 | | struct SearchByContentPrefix { |
79 | | prefix: String, |
80 | | } |
81 | | |
82 | | impl SchedulerIndexProvider for SearchByContentPrefix { |
83 | | const KEY_PREFIX: &'static str = "test:"; |
84 | | const INDEX_NAME: &'static str = "content_prefix"; |
85 | | type Versioned = TrueValue; |
86 | | |
87 | 0 | fn index_value(&self) -> Cow<'_, str> { |
88 | 0 | Cow::Borrowed(&self.prefix) |
89 | 0 | } |
90 | | } |
91 | | |
92 | | impl SchedulerStoreKeyProvider for SearchByContentPrefix { |
93 | | type Versioned = TrueValue; |
94 | | |
95 | 0 | fn get_key(&self) -> StoreKey<'static> { |
96 | 0 | StoreKey::Str(Cow::Owned("dummy_key".to_string())) |
97 | 0 | } |
98 | | } |
99 | | |
100 | | impl SchedulerStoreDecodeTo for SearchByContentPrefix { |
101 | | type DecodeOutput = TestSchedulerReturn; |
102 | | |
103 | 0 | fn decode(version: i64, data: Bytes) -> Result<Self::DecodeOutput, Error> { |
104 | 0 | TestSchedulerData::decode(version, data) |
105 | 0 | } |
106 | | } |
107 | | |
108 | | const MAX_KEY: u16 = 1024; |
109 | | |
110 | | /// Wrapper type for CLI parsing since we can't implement foreign traits on foreign types. |
111 | | #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, ValueEnum)] |
112 | | enum RedisModeArg { |
113 | | Cluster, |
114 | | Sentinel, |
115 | | #[default] |
116 | | Standard, |
117 | | } |
118 | | |
119 | | impl From<RedisModeArg> for RedisMode { |
120 | 0 | fn from(arg: RedisModeArg) -> Self { |
121 | 0 | match arg { |
122 | 0 | RedisModeArg::Standard => Self::Standard, |
123 | 0 | RedisModeArg::Sentinel => Self::Sentinel, |
124 | 0 | RedisModeArg::Cluster => Self::Cluster, |
125 | | } |
126 | 0 | } |
127 | | } |
128 | | |
129 | 0 | fn random_key() -> StoreKey<'static> { |
130 | 0 | let key = rand::rng().random_range(0..MAX_KEY); |
131 | 0 | StoreKey::new_str(&key.to_string()).into_owned() |
132 | 0 | } |
133 | | |
134 | | #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, ValueEnum)] |
135 | | enum TestMode { |
136 | | #[default] |
137 | | Random, |
138 | | Sequential, |
139 | | } |
140 | | |
141 | | #[derive(Parser, Debug)] |
142 | | #[command(version, about)] |
143 | | struct Args { |
144 | | #[arg(value_enum, short, long, default_value_t)] |
145 | | redis_mode: RedisModeArg, |
146 | | |
147 | | #[arg(value_enum, short, long, default_value_t)] |
148 | | mode: TestMode, |
149 | | } |
150 | | |
151 | 0 | async fn run<S: StoreDriver + SchedulerStore>( |
152 | 0 | store: Arc<S>, |
153 | 0 | max_loops: usize, |
154 | 0 | failed: Arc<RwLock<bool>>, |
155 | 0 | mode: TestMode, |
156 | 0 | ) -> Result<(), Error> { |
157 | 0 | let mut count = 0; |
158 | 0 | let in_flight = Arc::new(AtomicUsize::new(0)); |
159 | 0 | let fixed_action_value = if let Ok(str_action_value) = env::var("ACTION_VALUE") { |
160 | 0 | Some(str::parse::<usize>(&str_action_value).unwrap()) |
161 | | } else { |
162 | 0 | None |
163 | | }; |
164 | | |
165 | | loop { |
166 | 0 | if count % 1000 == 0 { |
167 | 0 | info!( |
168 | | "Loop count {count}. In flight: {}", |
169 | 0 | in_flight.load(Ordering::Relaxed) |
170 | | ); |
171 | 0 | if *failed.read().unwrap() { |
172 | 0 | return Err(Error::new( |
173 | 0 | Code::Internal, |
174 | 0 | "Failed in redis_store_tester".to_string(), |
175 | 0 | )); |
176 | 0 | } |
177 | 0 | } |
178 | 0 | if count == max_loops { |
179 | | loop { |
180 | 0 | let remaining = in_flight.load(Ordering::Relaxed); |
181 | 0 | if remaining == 0 { |
182 | 0 | return Ok(()); |
183 | 0 | } |
184 | 0 | info!(remaining, "Remaining"); |
185 | 0 | sleep(Duration::from_secs(1)).await; |
186 | | } |
187 | 0 | } |
188 | 0 | count += 1; |
189 | 0 | in_flight.fetch_add(1, Ordering::Relaxed); |
190 | | |
191 | 0 | let store_clone = store.clone(); |
192 | 0 | let local_fail = failed.clone(); |
193 | 0 | let local_in_flight = in_flight.clone(); |
194 | | |
195 | 0 | let max_action_value = 7; |
196 | 0 | let action_value = if let Some(av) = fixed_action_value { |
197 | 0 | av |
198 | | } else { |
199 | 0 | match mode { |
200 | 0 | TestMode::Random => rand::rng().random_range(0..max_action_value), |
201 | 0 | TestMode::Sequential => count % max_action_value, |
202 | | } |
203 | | }; |
204 | | |
205 | 0 | background_spawn!("action", async move { |
206 | 0 | async fn run_action<S: StoreDriver + SchedulerStore>( |
207 | 0 | action_value: usize, |
208 | 0 | store_clone: Arc<S>, |
209 | 0 | ) -> Result<(), Error> { |
210 | 0 | match action_value { |
211 | | 0 => { |
212 | 0 | store_clone.has(random_key()).await?; |
213 | | } |
214 | | 1 => { |
215 | 0 | let (mut tx, rx) = make_buf_channel_pair(); |
216 | 0 | tx.send(Bytes::from_static(b"12345")).await?; |
217 | 0 | tx.send_eof()?; |
218 | 0 | store_clone |
219 | 0 | .update(random_key(), rx, UploadSizeInfo::ExactSize(5)) |
220 | 0 | .await?; |
221 | | } |
222 | | 2 => { |
223 | 0 | let mut results = (0..MAX_KEY).map(|_| None).collect::<Vec<_>>(); |
224 | | |
225 | 0 | store_clone |
226 | 0 | .has_with_results( |
227 | 0 | &(0..MAX_KEY) |
228 | 0 | .map(|i| StoreKey::Str(Cow::Owned(i.to_string()))) |
229 | 0 | .collect::<Vec<_>>(), |
230 | 0 | &mut results, |
231 | | ) |
232 | 0 | .await?; |
233 | | } |
234 | | 3 => { |
235 | 0 | let key = random_key(); |
236 | 0 | store_clone |
237 | 0 | .update_oneshot(key.borrow(), Bytes::from_static(b"1234")) |
238 | 0 | .await?; |
239 | 0 | info!(?key, "Updated"); |
240 | | } |
241 | | 4 => { |
242 | 0 | let res = store_clone |
243 | 0 | .list(.., |_key| true) |
244 | 0 | .await |
245 | 0 | .err_tip(|| "In list")?; |
246 | 0 | info!(%res, "end list"); |
247 | | } |
248 | | 5 => { |
249 | 0 | let search_provider = SearchByContentPrefix { |
250 | 0 | prefix: "Searchable".to_string(), |
251 | 0 | }; |
252 | 0 | for i in 0..5 { |
253 | 0 | let data = TestSchedulerData { |
254 | 0 | key: format!("test:search_key_{i}"), |
255 | 0 | content: format!("Searchable content #{i}"), |
256 | 0 | version: 0, |
257 | 0 | }; |
258 | | |
259 | 0 | store_clone.update_data(data, None).await?; |
260 | | } |
261 | 0 | let search_results: Vec<_> = store_clone |
262 | 0 | .search_by_index_prefix(search_provider) |
263 | 0 | .await? |
264 | 0 | .try_collect() |
265 | 0 | .await?; |
266 | 0 | info!(?search_results, "search results"); |
267 | | } |
268 | | _ => { |
269 | 0 | let mut data = TestSchedulerData { |
270 | 0 | key: "test:scheduler_key_1".to_string(), |
271 | 0 | content: "Test scheduler data #1".to_string(), |
272 | 0 | version: 0, |
273 | 0 | }; |
274 | | |
275 | 0 | let res = store_clone.get_and_decode(data.clone()).await?; |
276 | 0 | if let Some(existing_data) = res { |
277 | 0 | data.version = existing_data.version + 1; |
278 | 0 | } |
279 | | |
280 | 0 | store_clone |
281 | 0 | .update_data(data, Some(Duration::from_mins(1))) |
282 | 0 | .await?; |
283 | | } |
284 | | } |
285 | 0 | Ok(()) |
286 | 0 | } |
287 | 0 | match run_action(action_value, store_clone).await { |
288 | 0 | Ok(()) => {} |
289 | 0 | Err(e) => { |
290 | 0 | error!(?e, "Error!"); |
291 | 0 | *local_fail.write().unwrap() = true; |
292 | | } |
293 | | } |
294 | 0 | local_in_flight.fetch_sub(1, Ordering::Relaxed); |
295 | 0 | }); |
296 | | } |
297 | 0 | } |
298 | | |
299 | | #[derive(Debug)] |
300 | | struct LoggingRemoveCallback {} |
301 | | |
302 | | impl RemoveItemCallback for LoggingRemoveCallback { |
303 | 0 | fn callback<'a>( |
304 | 0 | &'a self, |
305 | 0 | store_key: StoreKey<'a>, |
306 | 0 | ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> { |
307 | 0 | info!(?store_key, "Callback for removed item"); |
308 | 0 | Box::pin(async {}) |
309 | 0 | } |
310 | | } |
311 | | |
312 | 0 | fn main() -> Result<(), Box<dyn core::error::Error>> { |
313 | 0 | let args = Args::parse(); |
314 | 0 | let redis_mode: RedisMode = args.redis_mode.into(); |
315 | | |
316 | 0 | let failed = Arc::new(RwLock::new(false)); |
317 | 0 | let redis_host = env::var("REDIS_HOST").unwrap_or_else(|_| "127.0.0.1".to_string()); |
318 | 0 | let max_client_permits = env::var("MAX_REDIS_PERMITS") |
319 | 0 | .unwrap_or_else(|_| "100".to_string()) |
320 | 0 | .parse()?; |
321 | 0 | let max_loops: usize = env::var("MAX_LOOPS") |
322 | 0 | .unwrap_or_else(|_| "2000000".to_string()) |
323 | 0 | .parse()?; |
324 | | |
325 | | #[expect( |
326 | | clippy::disallowed_methods, |
327 | | reason = "`We need `tokio::runtime::Runtime::block_on` so we can get errors _after_ threads finished" |
328 | | )] |
329 | 0 | tokio::runtime::Builder::new_multi_thread() |
330 | 0 | .enable_all() |
331 | 0 | .build() |
332 | 0 | .unwrap() |
333 | 0 | .block_on(async { |
334 | | // The OTLP exporters need to run in a Tokio context. |
335 | 0 | spawn!("init tracing", async { init_tracing().await }) |
336 | 0 | .await? |
337 | 0 | .expect("Init tracing should work"); |
338 | | |
339 | 0 | let redis_port = match redis_mode { |
340 | 0 | RedisMode::Standard => 6379, |
341 | 0 | RedisMode::Sentinel => 26379, |
342 | 0 | RedisMode::Cluster => 7000, |
343 | | }; |
344 | 0 | let addr = match redis_mode { |
345 | 0 | RedisMode::Sentinel => format!("redis+sentinel://{redis_host}:{redis_port}/"), |
346 | 0 | _ => format!("redis://{redis_host}:{redis_port}/"), |
347 | | }; |
348 | 0 | let spec = RedisSpec { |
349 | 0 | addresses: vec![addr], |
350 | 0 | connection_timeout_ms: 1000, |
351 | 0 | max_client_permits, |
352 | 0 | mode: redis_mode, |
353 | 0 | ..Default::default() |
354 | 0 | }; |
355 | 0 | match spec.mode { |
356 | | RedisMode::Standard | RedisMode::Sentinel => { |
357 | 0 | let store = RedisStore::new_standard(spec).await?; |
358 | 0 | store |
359 | 0 | .clone() |
360 | 0 | .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; |
361 | 0 | run(store, max_loops, failed.clone(), args.mode).await |
362 | | } |
363 | | RedisMode::Cluster => { |
364 | 0 | let store = RedisStore::new_cluster(spec).await?; |
365 | 0 | store |
366 | 0 | .clone() |
367 | 0 | .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; |
368 | 0 | run(store, max_loops, failed.clone(), args.mode).await |
369 | | } |
370 | | } |
371 | 0 | }) |
372 | 0 | .unwrap(); |
373 | 0 | if *failed.read().unwrap() { |
374 | 0 | return Err(Error::new(Code::Internal, "Failed in redis_store_tester".to_string()).into()); |
375 | 0 | } |
376 | 0 | Ok(()) |
377 | 0 | } |