armature_cache/parallel.rs
1//! Parallel batch operations for cache stores.
2
3use crate::error::{CacheError, CacheResult};
4use crate::traits::CacheStore;
5use futures::StreamExt;
6use futures::future::join_all;
7use serde::{Serialize, de::DeserializeOwned};
8use std::collections::HashMap;
9use std::time::Duration;
10
11/// Default cap on how many warm-up factories [`ParallelCacheOps::warm_cache`]
12/// runs at once.
13///
14/// The factory is a user-supplied loader — typically a database query — so an
15/// unbounded fan-out would turn a 10,000-key warm-up into 10,000 simultaneous
16/// queries: precisely the stampede that
17/// [`CacheManager::get_or_set`](crate::manager::CacheManager::get_or_set)'s
18/// single-flight exists to prevent. Use
19/// [`ParallelCacheOps::warm_cache_with_concurrency`] to pick a different limit.
20pub const DEFAULT_WARM_CONCURRENCY: usize = 32;
21
22/// Parallel batch operations for cache stores.
23///
24/// This module provides high-performance batch operations that execute
25/// multiple cache operations concurrently, significantly reducing total latency.
26///
27/// # Performance
28///
29/// - **get_many**: 10-100x faster than sequential gets (depending on network latency)
30/// - **set_many**: 10-100x faster than sequential sets
31/// - **delete_many**: Similar performance gains
32///
33/// # Examples
34///
35/// ```no_run
36/// use armature_cache::*;
37/// use armature_cache::parallel::*;
38///
39/// # async fn example() -> CacheResult<()> {
40/// let cache = RedisCache::new(CacheConfig::redis("redis://localhost:6379")?).await?;
41///
42/// // Get multiple keys in parallel
43/// let keys = vec!["user:1", "user:2", "user:3"];
44/// let values = get_many_json(&cache, &keys).await?;
45///
46/// // Set multiple keys in parallel
47/// let items = vec![
48/// ("key1", "value1".to_string()),
49/// ("key2", "value2".to_string()),
50/// ];
51/// set_many_json(&cache, &items, None).await?;
52/// # Ok(())
53/// # }
54/// ```
55pub struct ParallelCacheOps;
56
57impl ParallelCacheOps {
58 /// Get multiple JSON values in parallel.
59 ///
60 /// # Arguments
61 ///
62 /// * `store` - The cache store
63 /// * `keys` - Slice of keys to fetch
64 ///
65 /// # Returns
66 ///
67 /// A vector of optional values in the same order as keys.
68 ///
69 /// # Examples
70 ///
71 /// ```no_run
72 /// use armature_cache::*;
73 /// use armature_cache::parallel::ParallelCacheOps;
74 ///
75 /// # async fn example() -> CacheResult<()> {
76 /// let cache = RedisCache::new(CacheConfig::redis("redis://localhost:6379")?).await?;
77 ///
78 /// let keys = vec!["key1", "key2", "key3"];
79 /// let values = ParallelCacheOps::get_many_json(&cache, &keys).await?;
80 ///
81 /// for (key, value) in keys.iter().zip(values.iter()) {
82 /// println!("{}: {:?}", key, value);
83 /// }
84 /// # Ok(())
85 /// # }
86 /// ```
87 pub async fn get_many_json<S: CacheStore>(
88 store: &S,
89 keys: &[&str],
90 ) -> CacheResult<Vec<Option<String>>> {
91 // Delegate to the store's native batch primitive. Backends like Redis
92 // collapse this into a single `MGET` round-trip; others fall back to the
93 // concurrent per-key loop. Order matches `keys` in both cases.
94 store.mget(keys).await
95 }
96
97 /// Get multiple typed values in parallel.
98 ///
99 /// # Type Parameters
100 ///
101 /// * `T` - The type to deserialize into
102 ///
103 /// # Examples
104 ///
105 /// ```no_run
106 /// use armature_cache::*;
107 /// use armature_cache::parallel::ParallelCacheOps;
108 /// use serde::{Deserialize, Serialize};
109 ///
110 /// #[derive(Serialize, Deserialize)]
111 /// struct User {
112 /// id: u64,
113 /// name: String,
114 /// }
115 ///
116 /// # async fn example() -> CacheResult<()> {
117 /// let cache = RedisCache::new(CacheConfig::redis("redis://localhost:6379")?).await?;
118 ///
119 /// let keys = vec!["user:1", "user:2", "user:3"];
120 /// let users: Vec<Option<User>> = ParallelCacheOps::get_many(&cache, &keys).await?;
121 /// # Ok(())
122 /// # }
123 /// ```
124 pub async fn get_many<S: CacheStore, T: DeserializeOwned>(
125 store: &S,
126 keys: &[&str],
127 ) -> CacheResult<Vec<Option<T>>> {
128 let json_values = Self::get_many_json(store, keys).await?;
129
130 json_values
131 .into_iter()
132 .map(|opt_json| {
133 opt_json
134 .map(|json| {
135 serde_json::from_str(&json)
136 .map_err(|e| CacheError::Deserialization(e.to_string()))
137 })
138 .transpose()
139 })
140 .collect()
141 }
142
143 /// Set multiple JSON values in parallel.
144 ///
145 /// # Arguments
146 ///
147 /// * `store` - The cache store
148 /// * `items` - Slice of (key, value) tuples
149 /// * `ttl` - Optional time-to-live for all items
150 ///
151 /// # Examples
152 ///
153 /// ```no_run
154 /// use armature_cache::*;
155 /// use armature_cache::parallel::ParallelCacheOps;
156 /// use std::time::Duration;
157 ///
158 /// # async fn example() -> CacheResult<()> {
159 /// let cache = RedisCache::new(CacheConfig::redis("redis://localhost:6379")?).await?;
160 ///
161 /// let items = vec![
162 /// ("key1", r#"{"value": 1}"#.to_string()),
163 /// ("key2", r#"{"value": 2}"#.to_string()),
164 /// ];
165 ///
166 /// ParallelCacheOps::set_many_json(&cache, &items, Some(Duration::from_secs(3600))).await?;
167 /// # Ok(())
168 /// # }
169 /// ```
170 pub async fn set_many_json<S: CacheStore>(
171 store: &S,
172 items: &[(&str, String)],
173 ttl: Option<Duration>,
174 ) -> CacheResult<()> {
175 // Delegate to the store's native batch primitive (e.g. Redis `MSET` /
176 // pipelined `SET ... EX`), falling back to the per-key loop otherwise.
177 store.mset(items, ttl).await
178 }
179
180 /// Set multiple typed values in parallel.
181 ///
182 /// # Type Parameters
183 ///
184 /// * `T` - The type to serialize from
185 ///
186 /// # Examples
187 ///
188 /// ```no_run
189 /// use armature_cache::*;
190 /// use armature_cache::parallel::ParallelCacheOps;
191 /// use serde::{Deserialize, Serialize};
192 ///
193 /// #[derive(Serialize, Deserialize)]
194 /// struct Counter {
195 /// count: u64,
196 /// }
197 ///
198 /// # async fn example() -> CacheResult<()> {
199 /// let cache = RedisCache::new(CacheConfig::redis("redis://localhost:6379")?).await?;
200 ///
201 /// let items = vec![
202 /// ("counter:1", Counter { count: 10 }),
203 /// ("counter:2", Counter { count: 20 }),
204 /// ];
205 ///
206 /// ParallelCacheOps::set_many(&cache, &items, None).await?;
207 /// # Ok(())
208 /// # }
209 /// ```
210 pub async fn set_many<S: CacheStore, T: Serialize>(
211 store: &S,
212 items: &[(&str, T)],
213 ttl: Option<Duration>,
214 ) -> CacheResult<()> {
215 let json_items: Result<Vec<_>, _> = items
216 .iter()
217 .map(|(key, value)| {
218 serde_json::to_string(value)
219 .map(|json| (*key, json))
220 .map_err(|e| CacheError::Serialization(e.to_string()))
221 })
222 .collect();
223
224 let json_items = json_items?;
225 let item_refs: Vec<_> = json_items.iter().map(|(k, v)| (*k, v.clone())).collect();
226
227 Self::set_many_json(store, &item_refs, ttl).await
228 }
229
230 /// Delete multiple keys in parallel.
231 ///
232 /// # Examples
233 ///
234 /// ```no_run
235 /// use armature_cache::*;
236 /// use armature_cache::parallel::ParallelCacheOps;
237 ///
238 /// # async fn example() -> CacheResult<()> {
239 /// let cache = RedisCache::new(CacheConfig::redis("redis://localhost:6379")?).await?;
240 ///
241 /// let keys = vec!["key1", "key2", "key3"];
242 /// ParallelCacheOps::delete_many(&cache, &keys).await?;
243 /// # Ok(())
244 /// # }
245 /// ```
246 pub async fn delete_many<S: CacheStore>(store: &S, keys: &[&str]) -> CacheResult<()> {
247 // Delegate to the store's native batch primitive (e.g. Redis variadic
248 // `DEL`), falling back to the concurrent per-key loop otherwise.
249 store.mdel(keys).await
250 }
251
252 /// Check if multiple keys exist in parallel.
253 ///
254 /// # Returns
255 ///
256 /// A vector of booleans indicating existence, in the same order as keys.
257 ///
258 /// # Examples
259 ///
260 /// ```no_run
261 /// use armature_cache::*;
262 /// use armature_cache::parallel::ParallelCacheOps;
263 ///
264 /// # async fn example() -> CacheResult<()> {
265 /// let cache = RedisCache::new(CacheConfig::redis("redis://localhost:6379")?).await?;
266 ///
267 /// let keys = vec!["key1", "key2", "key3"];
268 /// let exists = ParallelCacheOps::exists_many(&cache, &keys).await?;
269 ///
270 /// for (key, exists) in keys.iter().zip(exists.iter()) {
271 /// println!("{}: {}", key, exists);
272 /// }
273 /// # Ok(())
274 /// # }
275 /// ```
276 pub async fn exists_many<S: CacheStore>(store: &S, keys: &[&str]) -> CacheResult<Vec<bool>> {
277 let futures = keys.iter().map(|key| store.exists(key));
278 let results: Vec<CacheResult<bool>> = join_all(futures).await;
279
280 results.into_iter().collect()
281 }
282
283 /// Get TTL for multiple keys in parallel.
284 ///
285 /// # Returns
286 ///
287 /// A vector of optional durations, in the same order as keys.
288 ///
289 /// # Examples
290 ///
291 /// ```no_run
292 /// use armature_cache::*;
293 /// use armature_cache::parallel::ParallelCacheOps;
294 ///
295 /// # async fn example() -> CacheResult<()> {
296 /// let cache = RedisCache::new(CacheConfig::redis("redis://localhost:6379")?).await?;
297 ///
298 /// let keys = vec!["key1", "key2", "key3"];
299 /// let ttls = ParallelCacheOps::ttl_many(&cache, &keys).await?;
300 ///
301 /// for (key, ttl) in keys.iter().zip(ttls.iter()) {
302 /// println!("{}: {:?}", key, ttl);
303 /// }
304 /// # Ok(())
305 /// # }
306 /// ```
307 pub async fn ttl_many<S: CacheStore>(
308 store: &S,
309 keys: &[&str],
310 ) -> CacheResult<Vec<Option<Duration>>> {
311 let futures = keys.iter().map(|key| store.ttl(key));
312 let results: Vec<CacheResult<Option<Duration>>> = join_all(futures).await;
313
314 results.into_iter().collect()
315 }
316
317 /// Cache warming: preload multiple keys into cache.
318 ///
319 /// # Concurrency
320 ///
321 /// At most [`DEFAULT_WARM_CONCURRENCY`] factories run at a time. The
322 /// factory is the caller's loader (usually a DB fetch), so warming a large
323 /// key set without a bound would fire one query per key simultaneously.
324 /// Use [`Self::warm_cache_with_concurrency`] to choose the limit.
325 ///
326 /// # Type Parameters
327 ///
328 /// * `T` - The type to serialize
329 /// * `F` - Factory function that returns data for a given key
330 ///
331 /// # Examples
332 ///
333 /// ```ignore
334 /// use armature_cache::*;
335 /// use armature_cache::parallel::ParallelCacheOps;
336 /// use std::time::Duration;
337 ///
338 /// # async fn example() -> CacheResult<()> {
339 /// let cache = RedisCache::new(CacheConfig::redis("redis://localhost:6379")?).await?;
340 ///
341 /// let keys = vec!["user:1", "user:2", "user:3"];
342 ///
343 /// ParallelCacheOps::warm_cache(
344 /// &cache,
345 /// &keys,
346 /// Some(Duration::from_secs(3600)),
347 /// |key: &str| async move {
348 /// // Fetch from database
349 /// let data = format!("Data for {}", key);
350 /// Ok::<String, CacheError>(data)
351 /// },
352 /// ).await?;
353 /// # Ok(())
354 /// # }
355 /// ```
356 pub async fn warm_cache<S, T, F, Fut>(
357 store: &S,
358 keys: &[&str],
359 ttl: Option<Duration>,
360 factory: F,
361 ) -> CacheResult<()>
362 where
363 S: CacheStore,
364 T: Serialize,
365 F: Fn(&str) -> Fut,
366 Fut: std::future::Future<Output = CacheResult<T>>,
367 {
368 Self::warm_cache_with_concurrency(store, keys, ttl, DEFAULT_WARM_CONCURRENCY, factory).await
369 }
370
371 /// Cache warming with an explicit cap on concurrent factory invocations.
372 ///
373 /// `max_concurrent` is the number of factories (and their follow-up
374 /// writes) allowed to be in flight at once; `0` is treated as `1`. See
375 /// [`DEFAULT_WARM_CONCURRENCY`] for why the fan-out is bounded at all.
376 ///
377 /// The first error aborts the warm-up: keys already written stay written,
378 /// and the remaining ones are not attempted.
379 pub async fn warm_cache_with_concurrency<S, T, F, Fut>(
380 store: &S,
381 keys: &[&str],
382 ttl: Option<Duration>,
383 max_concurrent: usize,
384 factory: F,
385 ) -> CacheResult<()>
386 where
387 S: CacheStore,
388 T: Serialize,
389 F: Fn(&str) -> Fut,
390 Fut: std::future::Future<Output = CacheResult<T>>,
391 {
392 let limit = max_concurrent.max(1);
393 let factory = &factory;
394
395 let mut warmed = futures::stream::iter(keys.iter().map(|key| async move {
396 let value = factory(key).await?;
397 let json = serde_json::to_string(&value)
398 .map_err(|e| CacheError::Serialization(e.to_string()))?;
399 store.set_json(key, json, ttl).await?;
400 Ok::<(), CacheError>(())
401 }))
402 .buffer_unordered(limit);
403
404 while let Some(result) = warmed.next().await {
405 result?;
406 }
407
408 Ok(())
409 }
410}
411
412/// Helper functions for parallel cache operations.
413///
414/// These functions provide a more convenient API than `ParallelCacheOps` methods.
415/// Get multiple JSON values in parallel.
416pub async fn get_many_json<S: CacheStore>(
417 store: &S,
418 keys: &[&str],
419) -> CacheResult<Vec<Option<String>>> {
420 ParallelCacheOps::get_many_json(store, keys).await
421}
422
423/// Get multiple typed values in parallel.
424pub async fn get_many<S: CacheStore, T: DeserializeOwned>(
425 store: &S,
426 keys: &[&str],
427) -> CacheResult<Vec<Option<T>>> {
428 ParallelCacheOps::get_many(store, keys).await
429}
430
431/// Set multiple JSON values in parallel.
432pub async fn set_many_json<S: CacheStore>(
433 store: &S,
434 items: &[(&str, String)],
435 ttl: Option<Duration>,
436) -> CacheResult<()> {
437 ParallelCacheOps::set_many_json(store, items, ttl).await
438}
439
440/// Set multiple typed values in parallel.
441pub async fn set_many<S: CacheStore, T: Serialize>(
442 store: &S,
443 items: &[(&str, T)],
444 ttl: Option<Duration>,
445) -> CacheResult<()> {
446 ParallelCacheOps::set_many(store, items, ttl).await
447}
448
449/// Delete multiple keys in parallel.
450pub async fn delete_many<S: CacheStore>(store: &S, keys: &[&str]) -> CacheResult<()> {
451 ParallelCacheOps::delete_many(store, keys).await
452}
453
454/// Build a HashMap from multiple keys fetched in parallel.
455pub async fn get_many_as_map<S: CacheStore, T: DeserializeOwned>(
456 store: &S,
457 keys: &[&str],
458) -> CacheResult<HashMap<String, T>> {
459 let values = get_many(store, keys).await?;
460
461 let map: HashMap<String, T> = keys
462 .iter()
463 .zip(values)
464 .filter_map(|(key, opt_value)| opt_value.map(|value| (key.to_string(), value)))
465 .collect();
466
467 Ok(map)
468}
469
470#[cfg(test)]
471mod tests {
472 use super::*;
473 use crate::tiered::InMemoryCache;
474 use std::sync::Arc;
475 use std::sync::atomic::{AtomicUsize, Ordering};
476
477 /// Tracks how many factory invocations are in flight simultaneously.
478 #[derive(Default)]
479 struct ConcurrencyProbe {
480 current: AtomicUsize,
481 peak: AtomicUsize,
482 total: AtomicUsize,
483 }
484
485 impl ConcurrencyProbe {
486 fn enter(&self) {
487 let current = self.current.fetch_add(1, Ordering::SeqCst) + 1;
488 self.total.fetch_add(1, Ordering::SeqCst);
489 self.peak.fetch_max(current, Ordering::SeqCst);
490 }
491
492 fn leave(&self) {
493 self.current.fetch_sub(1, Ordering::SeqCst);
494 }
495 }
496
497 /// Regression: `warm_cache` used to `try_join_all` one future per key with
498 /// no limit, so warming N keys ran N user factories (documented as DB
499 /// fetches) simultaneously. The fan-out must now be bounded.
500 #[tokio::test]
501 async fn test_warm_cache_bounds_factory_concurrency() {
502 let cache = InMemoryCache::new();
503 let probe = Arc::new(ConcurrencyProbe::default());
504 let keys: Vec<String> = (0..64).map(|i| format!("k{i}")).collect();
505 let key_refs: Vec<&str> = keys.iter().map(|k| k.as_str()).collect();
506
507 const LIMIT: usize = 4;
508
509 ParallelCacheOps::warm_cache_with_concurrency(
510 &cache,
511 &key_refs,
512 None,
513 LIMIT,
514 |key: &str| {
515 let probe = probe.clone();
516 let key = key.to_string();
517 async move {
518 probe.enter();
519 // Suspend so the other buffered futures get a chance to run
520 // and the peak reflects real overlap.
521 for _ in 0..3 {
522 tokio::task::yield_now().await;
523 }
524 probe.leave();
525 Ok::<String, CacheError>(format!("value-for-{key}"))
526 }
527 },
528 )
529 .await
530 .unwrap();
531
532 assert_eq!(probe.total.load(Ordering::SeqCst), 64);
533 assert!(
534 probe.peak.load(Ordering::SeqCst) <= LIMIT,
535 "warm_cache ran {} factories at once, limit was {LIMIT}",
536 probe.peak.load(Ordering::SeqCst)
537 );
538
539 // Every key was still warmed.
540 assert_eq!(
541 cache.get_json("k0").await.unwrap(),
542 Some("\"value-for-k0\"".to_string())
543 );
544 assert_eq!(
545 cache.get_json("k63").await.unwrap(),
546 Some("\"value-for-k63\"".to_string())
547 );
548 }
549
550 /// The default entry point applies [`DEFAULT_WARM_CONCURRENCY`] rather
551 /// than fanning out over every key.
552 #[tokio::test]
553 async fn test_warm_cache_default_limit_applies() {
554 let cache = InMemoryCache::new();
555 let probe = Arc::new(ConcurrencyProbe::default());
556 let keys: Vec<String> = (0..DEFAULT_WARM_CONCURRENCY * 4)
557 .map(|i| format!("k{i}"))
558 .collect();
559 let key_refs: Vec<&str> = keys.iter().map(|k| k.as_str()).collect();
560
561 ParallelCacheOps::warm_cache(&cache, &key_refs, None, |_key: &str| {
562 let probe = probe.clone();
563 async move {
564 probe.enter();
565 for _ in 0..3 {
566 tokio::task::yield_now().await;
567 }
568 probe.leave();
569 Ok::<u32, CacheError>(1)
570 }
571 })
572 .await
573 .unwrap();
574
575 assert!(probe.peak.load(Ordering::SeqCst) <= DEFAULT_WARM_CONCURRENCY);
576 assert_eq!(
577 probe.total.load(Ordering::SeqCst),
578 DEFAULT_WARM_CONCURRENCY * 4
579 );
580 }
581
582 /// A concurrency limit of 0 is clamped to 1 rather than deadlocking or
583 /// silently doing nothing.
584 #[tokio::test]
585 async fn test_warm_cache_zero_concurrency_is_clamped_to_one() {
586 let cache = InMemoryCache::new();
587 let probe = Arc::new(ConcurrencyProbe::default());
588
589 ParallelCacheOps::warm_cache_with_concurrency(
590 &cache,
591 &["a", "b", "c"],
592 None,
593 0,
594 |_key: &str| {
595 let probe = probe.clone();
596 async move {
597 probe.enter();
598 tokio::task::yield_now().await;
599 probe.leave();
600 Ok::<u32, CacheError>(7)
601 }
602 },
603 )
604 .await
605 .unwrap();
606
607 assert_eq!(probe.peak.load(Ordering::SeqCst), 1);
608 assert_eq!(probe.total.load(Ordering::SeqCst), 3);
609 assert_eq!(cache.get_json("c").await.unwrap(), Some("7".to_string()));
610 }
611
612 /// A factory failure aborts the warm-up and surfaces the error.
613 #[tokio::test]
614 async fn test_warm_cache_propagates_factory_error() {
615 let cache = InMemoryCache::new();
616
617 let result = ParallelCacheOps::warm_cache_with_concurrency(
618 &cache,
619 &["a", "b"],
620 None,
621 1,
622 |key: &str| {
623 let fails = key == "b";
624 async move {
625 if fails {
626 Err(CacheError::Other("factory failed".to_string()))
627 } else {
628 Ok(1_u32)
629 }
630 }
631 },
632 )
633 .await;
634
635 assert!(result.is_err());
636 }
637}