shardline_server/
reconstruction_cache.rs

1use std::{
2    fmt::{Debug, Formatter, Result as FmtResult},
3    future::Future,
4    num::{NonZeroU64, NonZeroUsize},
5    sync::{
6        Arc,
7        atomic::{AtomicUsize, Ordering},
8    },
9    time::{Duration, Instant},
10};
11
12use serde_json::{from_slice, to_vec};
13use shardline_cache::{
14    AsyncReconstructionCache, DisabledReconstructionCache, MemoryReconstructionCache,
15    ReconstructionCacheKey, RedisReconstructionCache,
16};
17use shardline_protocol::RepositoryScope;
18
19use crate::{
20    FileReconstructionResponse, LocalBackend, ServerConfig, ServerConfigError, ServerError,
21};
22
23type SharedReconstructionCache = Arc<dyn AsyncReconstructionCache>;
24const MAX_RECONSTRUCTION_CACHE_PAYLOAD_BYTES: u64 = 67_108_864;
25
26/// Benchmarks one cold reconstruction load followed by one hot cache hit.
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub struct ReconstructionCacheBenchReport {
29    /// Cold reconstruction latency, including the backend load and cache insert.
30    pub cold_load_micros: u64,
31    /// Hot reconstruction latency served from cache.
32    pub hot_load_micros: u64,
33    /// Serialized reconstruction-response payload size used by the cache adapter.
34    pub response_bytes: u64,
35    /// Whether the second lookup avoided the backend loader.
36    pub cache_hit: bool,
37}
38
39/// Runtime reconstruction-cache service.
40#[derive(Clone)]
41pub struct ReconstructionCacheService {
42    adapter_name: &'static str,
43    adapter: SharedReconstructionCache,
44}
45
46impl Debug for ReconstructionCacheService {
47    fn fmt(&self, formatter: &mut Formatter<'_>) -> FmtResult {
48        formatter
49            .debug_struct("ReconstructionCacheService")
50            .field("adapter_name", &self.adapter_name)
51            .finish()
52    }
53}
54
55impl ReconstructionCacheService {
56    #[must_use]
57    pub fn disabled() -> Self {
58        Self {
59            adapter_name: ReconstructionCacheAdapter::Disabled.as_str(),
60            adapter: Arc::new(DisabledReconstructionCache::new()),
61        }
62    }
63
64    /// Builds the configured reconstruction-cache service.
65    ///
66    /// # Errors
67    ///
68    /// Returns [`ServerError`] when the configured adapter cannot initialize.
69    pub fn from_config(config: &ServerConfig) -> Result<Self, ServerError> {
70        match config.reconstruction_cache_adapter() {
71            ReconstructionCacheAdapter::Disabled => Ok(Self::disabled()),
72            ReconstructionCacheAdapter::Memory => Ok(Self {
73                adapter_name: ReconstructionCacheAdapter::Memory.as_str(),
74                adapter: Arc::new(MemoryReconstructionCache::new(
75                    config.reconstruction_cache_ttl_seconds(),
76                    config.reconstruction_cache_memory_max_entries(),
77                )),
78            }),
79            ReconstructionCacheAdapter::Redis => {
80                let redis_url = config
81                    .reconstruction_cache_redis_url()
82                    .ok_or(ServerError::MissingReconstructionCacheRedisUrl)?;
83                let adapter = RedisReconstructionCache::new(
84                    redis_url,
85                    config.reconstruction_cache_ttl_seconds(),
86                )?;
87                Ok(Self {
88                    adapter_name: ReconstructionCacheAdapter::Redis.as_str(),
89                    adapter: Arc::new(adapter),
90                })
91            }
92        }
93    }
94
95    fn for_tests(adapter_name: &'static str, adapter: SharedReconstructionCache) -> Self {
96        Self {
97            adapter_name,
98            adapter,
99        }
100    }
101
102    pub(crate) const fn backend_name(&self) -> &'static str {
103        self.adapter_name
104    }
105
106    pub(crate) async fn ready(&self) -> Result<(), ServerError> {
107        self.adapter.ready().await.map_err(ServerError::from)
108    }
109
110    pub(crate) async fn get_or_load<Load, LoadFuture>(
111        &self,
112        key: &ReconstructionCacheKey,
113        load: Load,
114    ) -> Result<FileReconstructionResponse, ServerError>
115    where
116        Load: FnOnce() -> LoadFuture,
117        LoadFuture: Future<Output = Result<FileReconstructionResponse, ServerError>>,
118    {
119        let cached = self.adapter.get(key).await;
120        if let Ok(Some(payload)) = cached
121            && payload_within_bound(&payload)
122        {
123            let parsed = from_slice::<FileReconstructionResponse>(&payload);
124            if let Ok(response) = parsed {
125                shardline_metrics::record_reconstruction_cache_hit();
126                return Ok(response);
127            }
128        }
129
130        shardline_metrics::record_reconstruction_cache_miss();
131        let response = load().await?;
132        let payload = to_vec(&response)?;
133        if payload_within_bound(&payload) {
134            let _ignored = self.adapter.put(key, &payload).await;
135        }
136        Ok(response)
137    }
138
139    pub(crate) fn version_key(
140        file_id: &str,
141        content_hash: &str,
142        repository_scope: Option<&RepositoryScope>,
143    ) -> ReconstructionCacheKey {
144        ReconstructionCacheKey::version(file_id, content_hash, repository_scope)
145    }
146}
147
148#[derive(Debug, Clone, Copy, PartialEq, Eq)]
149pub(crate) enum ReconstructionCacheAdapter {
150    Disabled,
151    Memory,
152    Redis,
153}
154
155impl ReconstructionCacheAdapter {
156    pub(crate) fn parse(value: &str) -> Result<Self, ServerConfigError> {
157        match value {
158            "disabled" => Ok(Self::Disabled),
159            "memory" => Ok(Self::Memory),
160            "redis" => Ok(Self::Redis),
161            _ => Err(ServerConfigError::InvalidReconstructionCacheAdapter),
162        }
163    }
164
165    pub(crate) const fn as_str(self) -> &'static str {
166        match self {
167            Self::Disabled => "disabled",
168            Self::Memory => "memory",
169            Self::Redis => "redis",
170        }
171    }
172}
173
174pub(crate) const DEFAULT_RECONSTRUCTION_CACHE_TTL_SECONDS: NonZeroU64 = match NonZeroU64::new(30) {
175    Some(value) => value,
176    None => NonZeroU64::MIN,
177};
178
179pub(crate) const DEFAULT_RECONSTRUCTION_CACHE_MEMORY_MAX_ENTRIES: NonZeroUsize =
180    match NonZeroUsize::new(4096) {
181        Some(value) => value,
182        None => NonZeroUsize::MIN,
183    };
184
185fn payload_within_bound(payload: &[u8]) -> bool {
186    let observed_bytes = u64::try_from(payload.len()).unwrap_or(u64::MAX);
187    observed_bytes <= MAX_RECONSTRUCTION_CACHE_PAYLOAD_BYTES
188}
189
190pub(crate) async fn benchmark_memory_reconstruction_cache_with_loader<Load, LoadFuture>(
191    file_id: &str,
192    content_hash: &str,
193    repository_scope: Option<&RepositoryScope>,
194    load: Load,
195) -> Result<ReconstructionCacheBenchReport, ServerError>
196where
197    Load: Fn() -> LoadFuture,
198    LoadFuture: Future<Output = Result<FileReconstructionResponse, ServerError>>,
199{
200    let ttl_seconds = NonZeroU64::new(60).unwrap_or(NonZeroU64::MIN);
201    let max_entries = NonZeroUsize::new(8).unwrap_or(NonZeroUsize::MIN);
202    let adapter: SharedReconstructionCache =
203        Arc::new(MemoryReconstructionCache::new(ttl_seconds, max_entries));
204    let cache = ReconstructionCacheService::for_tests("memory-bench", adapter);
205    let key = ReconstructionCacheService::version_key(file_id, content_hash, repository_scope);
206    let loader_calls = AtomicUsize::new(0);
207
208    let cold_started = Instant::now();
209    let cold = cache
210        .get_or_load(&key, || {
211            loader_calls.fetch_add(1, Ordering::SeqCst);
212            load()
213        })
214        .await?;
215    let cold_load_micros = duration_micros(cold_started.elapsed())?;
216
217    let hot_started = Instant::now();
218    let hot = cache
219        .get_or_load(&key, || {
220            loader_calls.fetch_add(1, Ordering::SeqCst);
221            load()
222        })
223        .await?;
224    let hot_load_micros = duration_micros(hot_started.elapsed())?;
225
226    debug_assert_eq!(cold, hot);
227    let response_bytes = u64::try_from(to_vec(&hot)?.len())?;
228
229    Ok(ReconstructionCacheBenchReport {
230        cold_load_micros,
231        hot_load_micros,
232        response_bytes,
233        cache_hit: loader_calls.load(Ordering::SeqCst) == 1,
234    })
235}
236
237/// Measures cold and hot reconstruction-cache behavior over the local backend path.
238///
239/// # Errors
240///
241/// Returns [`ServerError`] when reconstruction loading or cache serialization fails.
242pub async fn benchmark_memory_reconstruction_cache(
243    backend: &LocalBackend,
244    file_id: &str,
245    content_hash: &str,
246    repository_scope: Option<&RepositoryScope>,
247) -> Result<ReconstructionCacheBenchReport, ServerError> {
248    benchmark_memory_reconstruction_cache_with_loader(
249        file_id,
250        content_hash,
251        repository_scope,
252        || async move {
253            backend
254                .reconstruction(file_id, Some(content_hash), None, repository_scope)
255                .await
256        },
257    )
258    .await
259}
260
261fn duration_micros(duration: Duration) -> Result<u64, ServerError> {
262    u64::try_from(duration.as_micros()).map_err(ServerError::from)
263}
264
265#[cfg(test)]
266mod tests {
267    use std::{
268        num::{NonZeroU64, NonZeroUsize},
269        sync::{
270            Arc,
271            atomic::{AtomicUsize, Ordering},
272        },
273    };
274
275    use shardline_cache::{
276        AsyncReconstructionCache, MemoryReconstructionCache, ReconstructionCacheError,
277        ReconstructionCacheFuture, ReconstructionCacheKey,
278    };
279    use tokio::sync::Mutex;
280
281    use super::{
282        MAX_RECONSTRUCTION_CACHE_PAYLOAD_BYTES, ReconstructionCacheService,
283        SharedReconstructionCache,
284    };
285    use crate::FileReconstructionResponse;
286    use crate::xet_adapter::{
287        ReconstructionChunkRange, ReconstructionFetchInfo, ReconstructionTerm,
288        ReconstructionUrlRange,
289    };
290
291    #[derive(Debug)]
292    struct BrokenCache;
293
294    #[derive(Debug)]
295    struct StaticCache {
296        payload: Option<Vec<u8>>,
297        put_calls: Arc<AtomicUsize>,
298    }
299
300    impl AsyncReconstructionCache for BrokenCache {
301        fn ready(&self) -> ReconstructionCacheFuture<'_, ()> {
302            Box::pin(async { Err(ReconstructionCacheError::Operation) })
303        }
304
305        fn get<'operation>(
306            &'operation self,
307            _key: &'operation ReconstructionCacheKey,
308        ) -> ReconstructionCacheFuture<'operation, Option<Vec<u8>>> {
309            Box::pin(async { Err(ReconstructionCacheError::Operation) })
310        }
311
312        fn put<'operation>(
313            &'operation self,
314            _key: &'operation ReconstructionCacheKey,
315            _payload: &'operation [u8],
316        ) -> ReconstructionCacheFuture<'operation, ()> {
317            Box::pin(async { Err(ReconstructionCacheError::Operation) })
318        }
319
320        fn delete<'operation>(
321            &'operation self,
322            _key: &'operation ReconstructionCacheKey,
323        ) -> ReconstructionCacheFuture<'operation, bool> {
324            Box::pin(async { Err(ReconstructionCacheError::Operation) })
325        }
326    }
327
328    impl AsyncReconstructionCache for StaticCache {
329        fn ready(&self) -> ReconstructionCacheFuture<'_, ()> {
330            Box::pin(async { Ok(()) })
331        }
332
333        fn get<'operation>(
334            &'operation self,
335            _key: &'operation ReconstructionCacheKey,
336        ) -> ReconstructionCacheFuture<'operation, Option<Vec<u8>>> {
337            let payload = self.payload.clone();
338            Box::pin(async move { Ok(payload) })
339        }
340
341        fn put<'operation>(
342            &'operation self,
343            _key: &'operation ReconstructionCacheKey,
344            _payload: &'operation [u8],
345        ) -> ReconstructionCacheFuture<'operation, ()> {
346            self.put_calls.fetch_add(1, Ordering::SeqCst);
347            Box::pin(async { Ok(()) })
348        }
349
350        fn delete<'operation>(
351            &'operation self,
352            _key: &'operation ReconstructionCacheKey,
353        ) -> ReconstructionCacheFuture<'operation, bool> {
354            Box::pin(async { Ok(false) })
355        }
356    }
357
358    #[derive(Debug)]
359    struct CaptureCache {
360        put_calls: Arc<AtomicUsize>,
361        stored_payloads: Arc<Mutex<Vec<Vec<u8>>>>,
362    }
363
364    impl AsyncReconstructionCache for CaptureCache {
365        fn ready(&self) -> ReconstructionCacheFuture<'_, ()> {
366            Box::pin(async { Ok(()) })
367        }
368
369        fn get<'operation>(
370            &'operation self,
371            _key: &'operation ReconstructionCacheKey,
372        ) -> ReconstructionCacheFuture<'operation, Option<Vec<u8>>> {
373            Box::pin(async { Ok(None) })
374        }
375
376        fn put<'operation>(
377            &'operation self,
378            _key: &'operation ReconstructionCacheKey,
379            payload: &'operation [u8],
380        ) -> ReconstructionCacheFuture<'operation, ()> {
381            let payload = payload.to_vec();
382            let stored_payloads = Arc::clone(&self.stored_payloads);
383            self.put_calls.fetch_add(1, Ordering::SeqCst);
384            Box::pin(async move {
385                stored_payloads.lock().await.push(payload);
386                Ok(())
387            })
388        }
389
390        fn delete<'operation>(
391            &'operation self,
392            _key: &'operation ReconstructionCacheKey,
393        ) -> ReconstructionCacheFuture<'operation, bool> {
394            Box::pin(async { Ok(false) })
395        }
396    }
397
398    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
399    async fn cache_service_uses_cached_payload_after_first_load() {
400        let ttl_seconds = NonZeroU64::new(60).unwrap_or(NonZeroU64::MIN);
401        let adapter: SharedReconstructionCache = Arc::new(MemoryReconstructionCache::new(
402            ttl_seconds,
403            NonZeroUsize::MIN,
404        ));
405        let cache = ReconstructionCacheService::for_tests("memory", adapter);
406        let key = ReconstructionCacheKey::latest("asset.bin", None);
407        let loader_calls = AtomicUsize::new(0);
408
409        let first = cache
410            .get_or_load(&key, || {
411                loader_calls.fetch_add(1, Ordering::SeqCst);
412                async { Ok(sample_response("chunk-1")) }
413            })
414            .await;
415        let second = cache
416            .get_or_load(&key, || {
417                loader_calls.fetch_add(1, Ordering::SeqCst);
418                async { Ok(sample_response("chunk-2")) }
419            })
420            .await;
421
422        assert!(first.is_ok());
423        assert!(second.is_ok());
424        assert_eq!(loader_calls.load(Ordering::SeqCst), 1);
425        assert_eq!(
426            second
427                .ok()
428                .and_then(|response| response.terms.first().map(|term| term.hash.clone())),
429            Some("chunk-1".to_owned())
430        );
431    }
432
433    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
434    async fn cache_service_falls_back_to_loader_when_cache_adapter_errors() {
435        let adapter: SharedReconstructionCache = Arc::new(BrokenCache);
436        let cache = ReconstructionCacheService::for_tests("broken", adapter);
437        let key = ReconstructionCacheKey::latest("asset.bin", None);
438
439        let response = cache
440            .get_or_load(&key, || async { Ok(sample_response("chunk-1")) })
441            .await;
442
443        assert!(response.is_ok());
444        assert_eq!(
445            response
446                .ok()
447                .and_then(|value| value.terms.first().map(|term| term.unpacked_length)),
448            Some(4)
449        );
450    }
451
452    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
453    async fn cache_service_falls_back_to_loader_when_cached_payload_exceeds_bound() {
454        let put_calls = Arc::new(AtomicUsize::new(0));
455        let adapter: SharedReconstructionCache = Arc::new(StaticCache {
456            payload: Some(vec![
457                b'{';
458                MAX_RECONSTRUCTION_CACHE_PAYLOAD_BYTES as usize + 1
459            ]),
460            put_calls: Arc::clone(&put_calls),
461        });
462        let cache = ReconstructionCacheService::for_tests("static", adapter);
463        let key = ReconstructionCacheKey::latest("asset.bin", None);
464        let loader_calls = AtomicUsize::new(0);
465
466        let response = cache
467            .get_or_load(&key, || {
468                loader_calls.fetch_add(1, Ordering::SeqCst);
469                async { Ok(sample_response("chunk-1")) }
470            })
471            .await;
472
473        assert!(response.is_ok());
474        assert_eq!(loader_calls.load(Ordering::SeqCst), 1);
475        assert_eq!(put_calls.load(Ordering::SeqCst), 1);
476    }
477
478    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
479    async fn cache_service_skips_put_when_loaded_payload_exceeds_bound() {
480        let put_calls = Arc::new(AtomicUsize::new(0));
481        let stored_payloads = Arc::new(Mutex::new(Vec::new()));
482        let adapter: SharedReconstructionCache = Arc::new(CaptureCache {
483            put_calls: Arc::clone(&put_calls),
484            stored_payloads: Arc::clone(&stored_payloads),
485        });
486        let cache = ReconstructionCacheService::for_tests("capture", adapter);
487        let key = ReconstructionCacheKey::latest("asset.bin", None);
488        let oversized_hash =
489            "h".repeat(usize::try_from(MAX_RECONSTRUCTION_CACHE_PAYLOAD_BYTES).unwrap_or(0));
490
491        let response = cache
492            .get_or_load(&key, || async { Ok(sample_response(&oversized_hash)) })
493            .await;
494
495        assert!(response.is_ok());
496        assert_eq!(put_calls.load(Ordering::SeqCst), 0);
497        assert_eq!(stored_payloads.lock().await.len(), 0);
498    }
499
500    fn sample_response(hash: &str) -> FileReconstructionResponse {
501        FileReconstructionResponse {
502            offset_into_first_range: 0,
503            terms: vec![ReconstructionTerm {
504                hash: hash.to_owned(),
505                unpacked_length: 4,
506                range: ReconstructionChunkRange { start: 0, end: 1 },
507            }],
508            fetch_info: [(
509                hash.to_owned(),
510                vec![ReconstructionFetchInfo {
511                    range: ReconstructionChunkRange { start: 0, end: 1 },
512                    url: format!("https://cas.example.test/{hash}"),
513                    url_range: ReconstructionUrlRange { start: 0, end: 3 },
514                }],
515            )]
516            .into_iter()
517            .collect(),
518        }
519    }
520}
shardline_server/reconstruction_cache.rs

shardline_server/
reconstruction_cache.rs
