Coverage Report

Created: 2026-07-21 15:28

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/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
}