ursula_runtime/
cold_store.rs

1use std::collections::HashMap;
2use std::collections::VecDeque;
3use std::fmt;
4use std::future::Future;
5use std::io;
6use std::pin::Pin;
7use std::sync::Arc;
8use std::sync::Mutex;
9use std::sync::atomic::AtomicU64;
10use std::sync::atomic::Ordering;
11use std::time::Duration;
12#[cfg(not(madsim))]
13use std::time::SystemTime;
14#[cfg(not(madsim))]
15use std::time::UNIX_EPOCH;
16
17use bytes::Bytes;
18use crossbeam_utils::CachePadded;
19use opendal::Operator;
20use opendal::Scheme;
21use opendal::layers::RetryLayer;
22use opendal::layers::TimeoutLayer;
23use ursula_config::config::ColdBackend;
24use ursula_shard::BucketStreamId;
25use ursula_stream::ColdChunkRef;
26use ursula_stream::ObjectPayloadRef;
27
28use crate::ColdConfig;
29
30pub(crate) const DEFAULT_CONTENT_TYPE: &str = "application/octet-stream";
31// Keep this global atomic isolated from unrelated statics. This does not remove
32// contention on the counter itself, but avoids accidental false sharing with
33// adjacent data without adding a per-core sequence scheme to this low-frequency
34// object-key path.
35static COLD_CHUNK_SEQUENCE: CachePadded<AtomicU64> = CachePadded::new(AtomicU64::new(0));
36
37/// Wrap an S3 (opendal) operator with timeout and bounded-retry layers.
38///
39/// 1. **Per-attempt timeout** ([`TimeoutLayer`], inner): a blackholed endpoint
40///    — chaos `s3_unavailable`, or a "busy ESTAB" TCP socket whose future never
41///    gets polled again — otherwise hangs the caller until
42///    `net.ipv4.tcp_retries2` (~15 min). That is the original freeze: the raft
43///    state-machine worker awaits S3 inside `install_snapshot` (`&mut self`),
44///    and openraft type-level-serializes `apply` with it, so an unbounded S3
45///    stall freezes apply. Bounding every attempt keeps the worker progressing.
46/// 2. **Bounded retries** ([`RetryLayer`], outer): S3 answers `503 SlowDown`
47///    while a fresh key prefix warms up (and on transient network blips). These
48///    are `is_temporary()` errors; without retries a single 503 fails a
49///    snapshot upload/download, stalling a restarted node's rejoin/catch-up.
50///    Retries are bounded, so a sustained outage still fails fast enough (each
51///    attempt is timeout-bounded) and the cluster keeps progressing on quorum.
52pub(crate) fn with_s3_resilience(
53    operator: Operator,
54    timeout: Duration,
55    max_retries: usize,
56) -> Operator {
57    operator
58        .layer(
59            TimeoutLayer::new()
60                .with_timeout(timeout)
61                .with_io_timeout(timeout),
62        )
63        .layer(RetryLayer::new().with_max_times(max_retries).with_jitter())
64}
65
66#[derive(Clone)]
67pub struct ColdStore {
68    info: ColdStoreInfo,
69    operator: Operator,
70    read_cache: Option<Arc<ColdReadCache>>,
71    observer: Arc<Mutex<Option<ColdStoreObserver>>>,
72    fault_policy: Arc<Mutex<Option<ColdStoreFaultPolicy>>>,
73    delay_fn: Arc<Mutex<ColdStoreDelayFn>>,
74}
75
76pub type ColdStoreHandle = Arc<ColdStore>;
77
78type ColdStoreObserver = Arc<dyn Fn(ColdStoreEvent) + Send + Sync>;
79type ColdStoreFaultPolicy =
80    Arc<dyn Fn(&ColdStoreFaultContext) -> Option<ColdStoreFaultEffect> + Send + Sync>;
81type ColdStoreDelayFuture = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
82type ColdStoreDelayFn = Arc<dyn Fn(Duration) -> ColdStoreDelayFuture + Send + Sync>;
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum ColdStoreOperation {
86    WriteChunk,
87    DeleteChunk,
88    RemoveAll,
89    ReadObjectRange,
90}
91
92#[derive(Debug, Clone, PartialEq, Eq)]
93pub struct ColdStoreFaultContext {
94    pub operation: ColdStoreOperation,
95    pub stream_id: Option<BucketStreamId>,
96    pub path: String,
97    pub payload_len: Option<usize>,
98    pub read_start_offset: Option<u64>,
99    pub len: Option<usize>,
100    pub object_start: Option<u64>,
101    pub object_end: Option<u64>,
102    pub cached: Option<bool>,
103}
104
105#[derive(Debug, Clone, PartialEq, Eq)]
106pub struct ColdStoreFault {
107    pub message: String,
108}
109
110impl ColdStoreFault {
111    pub fn new(message: impl Into<String>) -> Self {
112        Self {
113            message: message.into(),
114        }
115    }
116}
117
118#[derive(Debug, Clone, Default, PartialEq, Eq)]
119pub struct ColdStoreFaultEffect {
120    pub delay: Option<Duration>,
121    pub error: Option<ColdStoreFault>,
122    pub truncate_read_to: Option<usize>,
123}
124
125impl ColdStoreFaultEffect {
126    pub fn delay(duration: Duration) -> Self {
127        Self {
128            delay: Some(duration),
129            error: None,
130            truncate_read_to: None,
131        }
132    }
133
134    pub fn fail(message: impl Into<String>) -> Self {
135        Self {
136            delay: None,
137            error: Some(ColdStoreFault::new(message)),
138            truncate_read_to: None,
139        }
140    }
141
142    pub fn delay_then_fail(duration: Duration, message: impl Into<String>) -> Self {
143        Self {
144            delay: Some(duration),
145            error: Some(ColdStoreFault::new(message)),
146            truncate_read_to: None,
147        }
148    }
149
150    pub fn truncate_read_to(len: usize) -> Self {
151        Self {
152            delay: None,
153            error: None,
154            truncate_read_to: Some(len),
155        }
156    }
157}
158
159#[derive(Debug, Clone, PartialEq, Eq)]
160pub enum ColdStoreEvent {
161    WriteChunkBegin {
162        path: String,
163        payload_len: usize,
164    },
165    WriteChunkComplete {
166        path: String,
167        object_size: u64,
168    },
169    DeleteChunkBegin {
170        path: String,
171    },
172    DeleteChunkComplete {
173        path: String,
174    },
175    RemoveAllBegin {
176        path: String,
177    },
178    RemoveAllComplete {
179        path: String,
180    },
181    ReadObjectRangeBegin {
182        stream_id: Option<BucketStreamId>,
183        path: String,
184        read_start_offset: u64,
185        len: usize,
186        object_start: u64,
187        object_end: u64,
188        cached: bool,
189    },
190    ReadObjectRangeComplete {
191        stream_id: Option<BucketStreamId>,
192        path: String,
193        read_start_offset: u64,
194        len: usize,
195        returned_len: usize,
196        cached: bool,
197    },
198    FaultInjected {
199        operation: ColdStoreOperation,
200        stream_id: Option<BucketStreamId>,
201        path: String,
202        message: String,
203    },
204    DelayInjected {
205        operation: ColdStoreOperation,
206        stream_id: Option<BucketStreamId>,
207        path: String,
208        delay_ms: u64,
209    },
210    TruncateInjected {
211        stream_id: Option<BucketStreamId>,
212        path: String,
213        requested_len: usize,
214        returned_len: usize,
215    },
216}
217
218impl fmt::Debug for ColdStore {
219    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
220        f.debug_struct("ColdStore")
221            .field("info", &self.info)
222            .field("operator", &self.operator)
223            .field("read_cache", &self.read_cache)
224            .finish_non_exhaustive()
225    }
226}
227
228#[derive(Clone, Debug, PartialEq, Eq)]
229pub struct ColdStoreInfo {
230    pub backend: &'static str,
231    pub root: Option<String>,
232    pub bucket: Option<String>,
233    pub region: Option<String>,
234    pub endpoint: Option<String>,
235}
236
237impl ColdStore {
238    pub fn memory() -> io::Result<Self> {
239        let operator = Operator::via_iter(Scheme::Memory, [])
240            .map_err(|err| io::Error::other(err.to_string()))?;
241        Ok(Self::from_operator(operator, ColdStoreInfo {
242            backend: "memory",
243            root: None,
244            bucket: None,
245            region: None,
246            endpoint: None,
247        }))
248    }
249
250    /// Build a [`ColdStore`] from an explicit [`ColdConfig`].
251    ///
252    /// The bootstrap layer assembles the typed config before calling this method;
253    /// this method is purely functional — it does not touch `std::env`.
254    ///
255    /// Returns `Err` when the backend is [`ColdBackend::None`];
256    /// the caller should skip construction when cold storage is disabled.
257    pub fn try_new(config: &ColdConfig) -> io::Result<Self> {
258        let mut store = match config.backend {
259            ColdBackend::None => {
260                return Err(io::Error::new(
261                    io::ErrorKind::InvalidInput,
262                    "ColdStore::try_new called with backend=none; \
263                     the caller should skip construction when cold storage is disabled",
264                ));
265            }
266            ColdBackend::Memory => Self::memory()?,
267            ColdBackend::S3 => Self::s3_from_config(config)?,
268        };
269        let cache = config.cache.clone().unwrap_or_default();
270        if cache.max_size.as_bytes() > 0 {
271            let cache_params = ColdReadCacheParams {
272                max_bytes: cache.max_size.as_bytes() as usize,
273                block_bytes: cache.block_size.as_bytes() as usize,
274                max_readahead_blocks: cache.readahead_blocks,
275            };
276            store = store.with_read_cache(cache_params);
277        }
278        Ok(store)
279    }
280
281    fn s3_from_config(config: &ColdConfig) -> io::Result<Self> {
282        let s3 = config.s3.as_ref().ok_or_else(|| {
283            io::Error::new(
284                io::ErrorKind::InvalidInput,
285                "s3 configuration is required when cold backend is s3",
286            )
287        })?;
288        let bucket = s3.bucket.as_deref().ok_or_else(|| {
289            io::Error::new(
290                io::ErrorKind::InvalidInput,
291                "s3 bucket is required when cold backend is s3",
292            )
293        })?;
294        if bucket.trim().is_empty() {
295            return Err(io::Error::new(
296                io::ErrorKind::InvalidInput,
297                "s3 bucket must not be empty",
298            ));
299        }
300
301        let mut builder = opendal::services::S3::default().bucket(bucket);
302        let mut configured_root = None;
303        if let Some(root) = config.root.as_deref()
304            && !root.trim().is_empty()
305        {
306            builder = builder.root(root);
307            configured_root = Some(root.to_owned());
308        }
309        let mut configured_region = None;
310        if let Some(region) = s3.region.as_deref()
311            && !region.trim().is_empty()
312        {
313            builder = builder.region(region);
314            configured_region = Some(region.to_owned());
315        }
316        let mut configured_endpoint = None;
317        if let Some(endpoint) = s3.endpoint.as_deref()
318            && !endpoint.trim().is_empty()
319        {
320            builder = builder.endpoint(endpoint);
321            configured_endpoint = Some(endpoint.to_owned());
322        }
323        if let Some(access_key_id) = s3.access_key_id.as_deref()
324            && !access_key_id.trim().is_empty()
325        {
326            builder = builder.access_key_id(access_key_id);
327        }
328        if let Some(secret_access_key) = s3.secret_access_key.as_deref()
329            && !secret_access_key.trim().is_empty()
330        {
331            builder = builder.secret_access_key(secret_access_key);
332        }
333        if let Some(session_token) = s3.session_token.as_deref()
334            && !session_token.trim().is_empty()
335        {
336            builder = builder.session_token(session_token);
337        }
338
339        Ok(Self::from_operator(
340            with_s3_resilience(
341                Operator::new(builder)
342                    .map_err(|err| io::Error::other(err.to_string()))?
343                    .finish(),
344                s3.timeout.as_duration(),
345                s3.max_retries,
346            ),
347            ColdStoreInfo {
348                backend: "s3",
349                root: configured_root,
350                bucket: Some(bucket.to_owned()),
351                region: configured_region,
352                endpoint: configured_endpoint,
353            },
354        ))
355    }
356
357    fn from_operator(operator: Operator, info: ColdStoreInfo) -> Self {
358        Self {
359            info,
360            operator,
361            read_cache: None,
362            observer: Arc::new(Mutex::new(None)),
363            fault_policy: Arc::new(Mutex::new(None)),
364            delay_fn: Arc::new(Mutex::new(default_cold_store_delay_fn())),
365        }
366    }
367
368    pub fn info(&self) -> &ColdStoreInfo {
369        &self.info
370    }
371
372    pub fn with_read_cache(mut self, config: ColdReadCacheParams) -> Self {
373        self.read_cache = Some(Arc::new(ColdReadCache::new(config)));
374        self
375    }
376
377    pub fn without_read_cache(mut self) -> Self {
378        self.read_cache = None;
379        self
380    }
381
382    pub fn set_observer(&self, observer: impl Fn(ColdStoreEvent) + Send + Sync + 'static) {
383        *self.observer.lock().expect("cold store observer mutex") = Some(Arc::new(observer));
384    }
385
386    pub fn set_fault_policy(
387        &self,
388        policy: impl Fn(&ColdStoreFaultContext) -> Option<ColdStoreFaultEffect> + Send + Sync + 'static,
389    ) {
390        *self
391            .fault_policy
392            .lock()
393            .expect("cold store fault policy mutex") = Some(Arc::new(policy));
394    }
395
396    pub fn clear_fault_policy(&self) {
397        *self
398            .fault_policy
399            .lock()
400            .expect("cold store fault policy mutex") = None;
401    }
402
403    pub fn set_delay_fn<F, Fut>(&self, delay_fn: F)
404    where
405        F: Fn(Duration) -> Fut + Send + Sync + 'static,
406        Fut: Future<Output = ()> + Send + 'static,
407    {
408        *self.delay_fn.lock().expect("cold store delay fn mutex") =
409            Arc::new(move |duration| Box::pin(delay_fn(duration)));
410    }
411
412    #[cfg(test)]
413    pub(crate) fn cached_block_count(&self) -> usize {
414        self.read_cache
415            .as_ref()
416            .map(|cache| cache.block_count())
417            .unwrap_or(0)
418    }
419
420    pub async fn write_chunk(&self, path: &str, payload: &[u8]) -> io::Result<u64> {
421        if path.trim().is_empty() {
422            return Err(io::Error::new(
423                io::ErrorKind::InvalidInput,
424                "cold chunk path must not be empty",
425            ));
426        }
427        self.notify(ColdStoreEvent::WriteChunkBegin {
428            path: path.to_owned(),
429            payload_len: payload.len(),
430        });
431        let _applied_fault = self
432            .maybe_apply_fault_effect(ColdStoreFaultContext {
433                operation: ColdStoreOperation::WriteChunk,
434                stream_id: None,
435                path: path.to_owned(),
436                payload_len: Some(payload.len()),
437                read_start_offset: None,
438                len: None,
439                object_start: None,
440                object_end: None,
441                cached: None,
442            })
443            .await?;
444        self.operator
445            .write(path, payload.to_vec())
446            .await
447            .map_err(|err| cold_store_io_error(path, err))?;
448        let object_size = u64::try_from(payload.len()).expect("payload len fits u64");
449        self.notify(ColdStoreEvent::WriteChunkComplete {
450            path: path.to_owned(),
451            object_size,
452        });
453        Ok(object_size)
454    }
455
456    pub(crate) async fn write_cold_index_page(
457        &self,
458        path: &str,
459        payload: &[u8],
460    ) -> io::Result<u64> {
461        if path.trim().is_empty() {
462            return Err(io::Error::new(
463                io::ErrorKind::InvalidInput,
464                "cold index page path must not be empty",
465            ));
466        }
467        self.operator
468            .write(path, payload.to_vec())
469            .await
470            .map_err(|err| cold_store_io_error(path, err))?;
471        Ok(u64::try_from(payload.len()).expect("payload len fits u64"))
472    }
473
474    #[tracing::instrument(name = "cold.read_index", level = "debug", skip_all)]
475    pub(crate) async fn read_cold_index_page(&self, path: &str) -> io::Result<Option<Vec<u8>>> {
476        if path.trim().is_empty() {
477            return Err(io::Error::new(
478                io::ErrorKind::InvalidInput,
479                "cold index page path must not be empty",
480            ));
481        }
482        match self.operator.read(path).await {
483            Ok(bytes) => Ok(Some(bytes.to_bytes().to_vec())),
484            Err(err) if err.kind() == opendal::ErrorKind::NotFound => Ok(None),
485            Err(err) => Err(cold_store_io_error(path, err)),
486        }
487    }
488
489    pub async fn delete_chunk(&self, path: &str) -> io::Result<()> {
490        if path.trim().is_empty() {
491            return Err(io::Error::new(
492                io::ErrorKind::InvalidInput,
493                "cold chunk path must not be empty",
494            ));
495        }
496        self.notify(ColdStoreEvent::DeleteChunkBegin {
497            path: path.to_owned(),
498        });
499        let _applied_fault = self
500            .maybe_apply_fault_effect(ColdStoreFaultContext {
501                operation: ColdStoreOperation::DeleteChunk,
502                stream_id: None,
503                path: path.to_owned(),
504                payload_len: None,
505                read_start_offset: None,
506                len: None,
507                object_start: None,
508                object_end: None,
509                cached: None,
510            })
511            .await?;
512        self.operator
513            .delete(path)
514            .await
515            .map_err(|err| cold_store_io_error(path, err))?;
516        if let Some(cache) = &self.read_cache {
517            cache.invalidate_path(path);
518        }
519        self.notify(ColdStoreEvent::DeleteChunkComplete {
520            path: path.to_owned(),
521        });
522        Ok(())
523    }
524
525    pub async fn remove_all(&self, path: &str) -> io::Result<()> {
526        self.notify(ColdStoreEvent::RemoveAllBegin {
527            path: path.to_owned(),
528        });
529        let _applied_fault = self
530            .maybe_apply_fault_effect(ColdStoreFaultContext {
531                operation: ColdStoreOperation::RemoveAll,
532                stream_id: None,
533                path: path.to_owned(),
534                payload_len: None,
535                read_start_offset: None,
536                len: None,
537                object_start: None,
538                object_end: None,
539                cached: None,
540            })
541            .await?;
542        self.operator
543            .remove_all(path)
544            .await
545            .map_err(|err| cold_store_io_error(path, err))?;
546        if let Some(cache) = &self.read_cache {
547            cache.invalidate_prefix(path);
548        }
549        self.notify(ColdStoreEvent::RemoveAllComplete {
550            path: path.to_owned(),
551        });
552        Ok(())
553    }
554
555    pub async fn read_chunk_range(
556        &self,
557        chunk: &ColdChunkRef,
558        read_start_offset: u64,
559        len: usize,
560    ) -> io::Result<Vec<u8>> {
561        let object = ObjectPayloadRef {
562            start_offset: chunk.start_offset,
563            end_offset: chunk.end_offset,
564            s3_path: chunk.s3_path.clone(),
565            object_size: chunk.object_size,
566        };
567        self.read_object_range(&object, read_start_offset, len)
568            .await
569    }
570
571    pub async fn read_object_range_for_stream(
572        &self,
573        stream_id: &BucketStreamId,
574        object: &ObjectPayloadRef,
575        read_start_offset: u64,
576        len: usize,
577    ) -> io::Result<Vec<u8>> {
578        self.read_object_range_inner(Some(stream_id), object, read_start_offset, len)
579            .await
580    }
581
582    pub async fn read_object_range(
583        &self,
584        object: &ObjectPayloadRef,
585        read_start_offset: u64,
586        len: usize,
587    ) -> io::Result<Vec<u8>> {
588        self.read_object_range_inner(None, object, read_start_offset, len)
589            .await
590    }
591
592    #[tracing::instrument(
593        name = "cold.read_chunk",
594        level = "debug",
595        skip_all,
596        fields(
597            start_offset = object.start_offset,
598            end_offset = object.end_offset,
599            object_size = object.object_size,
600            len = len,
601        ),
602    )]
603    async fn read_object_range_inner(
604        &self,
605        stream_id: Option<&BucketStreamId>,
606        object: &ObjectPayloadRef,
607        read_start_offset: u64,
608        len: usize,
609    ) -> io::Result<Vec<u8>> {
610        if len == 0 {
611            return Ok(Vec::new());
612        }
613        let len_u64 = u64::try_from(len).map_err(|_| {
614            io::Error::new(io::ErrorKind::InvalidInput, "cold read length exceeds u64")
615        })?;
616        let read_end = read_start_offset.checked_add(len_u64).ok_or_else(|| {
617            io::Error::new(io::ErrorKind::InvalidInput, "cold read range overflow")
618        })?;
619        if read_start_offset < object.start_offset || read_end > object.end_offset {
620            return Err(io::Error::new(
621                io::ErrorKind::InvalidInput,
622                format!(
623                    "cold read range [{read_start_offset}..{read_end}) is outside object segment [{}..{})",
624                    object.start_offset, object.end_offset
625                ),
626            ));
627        }
628        let object_start = read_start_offset - object.start_offset;
629        let object_end = object_start.checked_add(len_u64).ok_or_else(|| {
630            io::Error::new(io::ErrorKind::InvalidInput, "cold read range overflow")
631        })?;
632        if object_end > object.object_size {
633            return Err(io::Error::new(
634                io::ErrorKind::InvalidData,
635                format!(
636                    "cold read range [{object_start}..{object_end}) is outside object '{}' size {}",
637                    object.s3_path, object.object_size
638                ),
639            ));
640        }
641        let cached = self.read_cache.is_some();
642        self.notify(ColdStoreEvent::ReadObjectRangeBegin {
643            stream_id: stream_id.cloned(),
644            path: object.s3_path.clone(),
645            read_start_offset,
646            len,
647            object_start,
648            object_end,
649            cached,
650        });
651        let applied_fault = self
652            .maybe_apply_fault_effect(ColdStoreFaultContext {
653                operation: ColdStoreOperation::ReadObjectRange,
654                stream_id: stream_id.cloned(),
655                path: object.s3_path.clone(),
656                payload_len: None,
657                read_start_offset: Some(read_start_offset),
658                len: Some(len),
659                object_start: Some(object_start),
660                object_end: Some(object_end),
661                cached: Some(cached),
662            })
663            .await?;
664        let mut bytes = if let Some(cache) = &self.read_cache {
665            let bytes = self
666                .read_object_range_cached(cache, object, object_start, object_end, len)
667                .await?;
668            if let Some(stream_id) = stream_id {
669                let readahead_blocks = cache.record_stream_read(stream_id, read_start_offset, len);
670                if readahead_blocks > 0 {
671                    self.spawn_readahead(object.clone(), object_end, readahead_blocks);
672                }
673            }
674            bytes
675        } else {
676            self.read_object_range_uncached(object, object_start, object_end, len)
677                .await?
678        };
679        if let Some(returned_len) = applied_fault.truncate_read_to {
680            let returned_len = returned_len.min(bytes.len());
681            bytes.truncate(returned_len);
682            self.notify(ColdStoreEvent::TruncateInjected {
683                stream_id: stream_id.cloned(),
684                path: object.s3_path.clone(),
685                requested_len: len,
686                returned_len,
687            });
688        }
689        if bytes.len() != len {
690            return Err(io::Error::new(
691                io::ErrorKind::InvalidData,
692                format!(
693                    "cold object '{}' returned {} bytes for requested range [{}..{})",
694                    object.s3_path,
695                    bytes.len(),
696                    object_start,
697                    object_end
698                ),
699            ));
700        }
701        self.notify(ColdStoreEvent::ReadObjectRangeComplete {
702            stream_id: stream_id.cloned(),
703            path: object.s3_path.clone(),
704            read_start_offset,
705            len,
706            returned_len: bytes.len(),
707            cached,
708        });
709        Ok(bytes)
710    }
711
712    async fn read_object_range_uncached(
713        &self,
714        object: &ObjectPayloadRef,
715        object_start: u64,
716        object_end: u64,
717        len: usize,
718    ) -> io::Result<Vec<u8>> {
719        let bytes = self
720            .operator
721            .read_with(&object.s3_path)
722            .range(object_start..object_end)
723            .await
724            .map_err(|err| cold_store_io_error(&object.s3_path, err))?
725            .to_bytes();
726        if bytes.len() != len {
727            return Err(io::Error::new(
728                io::ErrorKind::InvalidData,
729                format!(
730                    "cold object '{}' returned {} bytes for requested range [{}..{})",
731                    object.s3_path,
732                    bytes.len(),
733                    object_start,
734                    object_end
735                ),
736            ));
737        }
738        Ok(bytes.to_vec())
739    }
740
741    async fn read_object_range_cached(
742        &self,
743        cache: &ColdReadCache,
744        object: &ObjectPayloadRef,
745        object_start: u64,
746        object_end: u64,
747        len: usize,
748    ) -> io::Result<Vec<u8>> {
749        let mut payload = Vec::with_capacity(len);
750        let block_size = cache.block_size();
751        let first_block = object_start / block_size;
752        let last_block = (object_end - 1) / block_size;
753        for block_index in first_block..=last_block {
754            let block_start = block_index * block_size;
755            let block_end = block_start
756                .saturating_add(block_size)
757                .min(object.object_size);
758            let block = self
759                .read_cached_block(
760                    cache,
761                    object.s3_path.clone(),
762                    object.object_size,
763                    block_index,
764                    block_start,
765                    block_end,
766                )
767                .await?;
768            let slice_start = usize::try_from(object_start.max(block_start) - block_start)
769                .expect("cache slice start fits usize");
770            let slice_end = usize::try_from(object_end.min(block_end) - block_start)
771                .expect("cache slice end fits usize");
772            payload.extend_from_slice(&block.slice(slice_start..slice_end));
773        }
774        if payload.len() != len {
775            return Err(io::Error::new(
776                io::ErrorKind::InvalidData,
777                format!(
778                    "cold object '{}' returned {} bytes for requested range [{}..{})",
779                    object.s3_path,
780                    payload.len(),
781                    object_start,
782                    object_end
783                ),
784            ));
785        }
786        Ok(payload)
787    }
788
789    async fn read_cached_block(
790        &self,
791        cache: &ColdReadCache,
792        path: String,
793        object_size: u64,
794        block_index: u64,
795        block_start: u64,
796        block_end: u64,
797    ) -> io::Result<Bytes> {
798        if let Some(bytes) = cache.get(&path, block_index) {
799            return Ok(bytes);
800        }
801        let bytes = self
802            .operator
803            .read_with(&path)
804            .range(block_start..block_end)
805            .await
806            .map_err(|err| cold_store_io_error(&path, err))?
807            .to_bytes();
808        let expected_len = usize::try_from(block_end - block_start).map_err(|_| {
809            io::Error::new(
810                io::ErrorKind::InvalidData,
811                "cold cache block length exceeds usize",
812            )
813        })?;
814        if bytes.len() != expected_len {
815            return Err(io::Error::new(
816                io::ErrorKind::InvalidData,
817                format!(
818                    "cold object '{path}' returned {} bytes for cache block [{}..{}) of object size {object_size}",
819                    bytes.len(),
820                    block_start,
821                    block_end
822                ),
823            ));
824        }
825        cache.insert(path, block_index, bytes.clone());
826        Ok(bytes)
827    }
828
829    fn spawn_readahead(&self, object: ObjectPayloadRef, object_end: u64, readahead_blocks: usize) {
830        let Some(cache) = self.read_cache.clone() else {
831            return;
832        };
833        let block_size = cache.block_size();
834        let mut block_index = object_end.div_ceil(block_size);
835        let store = self.clone();
836        crate::rt::spawn(async move {
837            for _ in 0..readahead_blocks {
838                let block_start = block_index * block_size;
839                if block_start >= object.object_size {
840                    break;
841                }
842                let block_end = block_start
843                    .saturating_add(block_size)
844                    .min(object.object_size);
845                if cache.get(&object.s3_path, block_index).is_none() {
846                    let _ = store
847                        .read_cached_block(
848                            &cache,
849                            object.s3_path.clone(),
850                            object.object_size,
851                            block_index,
852                            block_start,
853                            block_end,
854                        )
855                        .await;
856                }
857                block_index += 1;
858            }
859        });
860    }
861
862    fn notify(&self, event: ColdStoreEvent) {
863        let observer = self
864            .observer
865            .lock()
866            .expect("cold store observer mutex")
867            .clone();
868        if let Some(observer) = observer {
869            observer(event);
870        }
871    }
872
873    async fn maybe_apply_fault_effect(
874        &self,
875        context: ColdStoreFaultContext,
876    ) -> io::Result<ColdStoreAppliedFault> {
877        let policy = self
878            .fault_policy
879            .lock()
880            .expect("cold store fault policy mutex")
881            .clone();
882        let Some(policy) = policy else {
883            return Ok(ColdStoreAppliedFault::default());
884        };
885        let Some(effect) = policy(&context) else {
886            return Ok(ColdStoreAppliedFault::default());
887        };
888        if let Some(delay) = effect.delay {
889            self.notify(ColdStoreEvent::DelayInjected {
890                operation: context.operation,
891                stream_id: context.stream_id.clone(),
892                path: context.path.clone(),
893                delay_ms: duration_ms(delay),
894            });
895            let delay_fn = self
896                .delay_fn
897                .lock()
898                .expect("cold store delay fn mutex")
899                .clone();
900            delay_fn(delay).await;
901        }
902        if let Some(fault) = effect.error {
903            self.notify(ColdStoreEvent::FaultInjected {
904                operation: context.operation,
905                stream_id: context.stream_id,
906                path: context.path.clone(),
907                message: fault.message.clone(),
908            });
909            return Err(io::Error::other(format!(
910                "cold store fault injected for {} '{}': {}",
911                context.operation.as_str(),
912                context.path,
913                fault.message
914            )));
915        }
916        Ok(ColdStoreAppliedFault {
917            truncate_read_to: effect.truncate_read_to,
918        })
919    }
920}
921
922#[derive(Debug, Default)]
923struct ColdStoreAppliedFault {
924    truncate_read_to: Option<usize>,
925}
926
927impl ColdStoreOperation {
928    fn as_str(self) -> &'static str {
929        match self {
930            Self::WriteChunk => "write_chunk",
931            Self::DeleteChunk => "delete_chunk",
932            Self::RemoveAll => "remove_all",
933            Self::ReadObjectRange => "read_object_range",
934        }
935    }
936}
937
938fn default_cold_store_delay_fn() -> ColdStoreDelayFn {
939    Arc::new(|duration| Box::pin(crate::rt::time::sleep(duration)))
940}
941
942fn duration_ms(duration: Duration) -> u64 {
943    u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
944}
945
946/// Runtime parameters for the optional cold-read cache.
947///
948/// Kept separate from the serde [`ursula_config::ColdCacheConfig`] because
949/// `ColdReadCache` accesses these fields on the hot path.
950#[derive(Debug, Clone, Copy)]
951pub struct ColdReadCacheParams {
952    pub max_bytes: usize,
953    pub block_bytes: usize,
954    pub max_readahead_blocks: usize,
955}
956
957#[derive(Debug)]
958struct ColdReadCache {
959    config: ColdReadCacheParams,
960    inner: Mutex<ColdReadCacheInner>,
961}
962
963#[derive(Debug, Default)]
964struct ColdReadCacheInner {
965    blocks: HashMap<ColdCacheKey, ColdCacheEntry>,
966    lru: VecDeque<(ColdCacheKey, u64)>,
967    current_bytes: usize,
968    generation: u64,
969    readers: HashMap<BucketStreamId, StreamReadState>,
970}
971
972#[derive(Debug, Clone, PartialEq, Eq, Hash)]
973struct ColdCacheKey {
974    path: String,
975    block_index: u64,
976}
977
978#[derive(Debug)]
979struct ColdCacheEntry {
980    bytes: Bytes,
981    generation: u64,
982}
983
984#[derive(Debug, Default)]
985struct StreamReadState {
986    next_offset: u64,
987    sequential_score: usize,
988}
989
990impl ColdReadCache {
991    fn new(config: ColdReadCacheParams) -> Self {
992        let block_bytes = config.block_bytes.max(1);
993        Self {
994            config: ColdReadCacheParams {
995                max_bytes: config.max_bytes,
996                block_bytes,
997                max_readahead_blocks: config.max_readahead_blocks,
998            },
999            inner: Mutex::new(ColdReadCacheInner::default()),
1000        }
1001    }
1002
1003    fn block_size(&self) -> u64 {
1004        u64::try_from(self.config.block_bytes.max(1)).expect("cache block size fits u64")
1005    }
1006
1007    fn get(&self, path: &str, block_index: u64) -> Option<Bytes> {
1008        let mut inner = self.inner.lock().expect("cold cache mutex poisoned");
1009        let key = ColdCacheKey {
1010            path: path.to_owned(),
1011            block_index,
1012        };
1013        let bytes = inner.blocks.get(&key)?.bytes.clone();
1014        Self::touch(&mut inner, key);
1015        Some(bytes)
1016    }
1017
1018    fn insert(&self, path: String, block_index: u64, bytes: Bytes) {
1019        if bytes.len() > self.config.max_bytes || self.config.max_bytes == 0 {
1020            return;
1021        }
1022        let mut inner = self.inner.lock().expect("cold cache mutex poisoned");
1023        let key = ColdCacheKey { path, block_index };
1024        if let Some(previous) = inner.blocks.remove(&key) {
1025            inner.current_bytes = inner.current_bytes.saturating_sub(previous.bytes.len());
1026        }
1027        let generation = Self::next_generation(&mut inner);
1028        inner.current_bytes = inner.current_bytes.saturating_add(bytes.len());
1029        inner
1030            .blocks
1031            .insert(key.clone(), ColdCacheEntry { bytes, generation });
1032        inner.lru.push_back((key, generation));
1033        self.evict_locked(&mut inner);
1034        Self::compact_lru_if_needed(&mut inner);
1035    }
1036
1037    fn record_stream_read(
1038        &self,
1039        stream_id: &BucketStreamId,
1040        read_start_offset: u64,
1041        len: usize,
1042    ) -> usize {
1043        let mut inner = self.inner.lock().expect("cold cache mutex poisoned");
1044        let state = inner.readers.entry(stream_id.clone()).or_default();
1045        if read_start_offset == state.next_offset {
1046            state.sequential_score = state
1047                .sequential_score
1048                .saturating_add(1)
1049                .min(self.config.max_readahead_blocks);
1050        } else {
1051            state.sequential_score = 0;
1052        }
1053        state.next_offset =
1054            read_start_offset.saturating_add(u64::try_from(len).unwrap_or(u64::MAX));
1055        state.sequential_score.min(self.config.max_readahead_blocks)
1056    }
1057
1058    fn invalidate_path(&self, path: &str) {
1059        let mut inner = self.inner.lock().expect("cold cache mutex poisoned");
1060        let keys = inner
1061            .blocks
1062            .keys()
1063            .filter(|key| key.path == path)
1064            .cloned()
1065            .collect::<Vec<_>>();
1066        for key in keys {
1067            if let Some(entry) = inner.blocks.remove(&key) {
1068                inner.current_bytes = inner.current_bytes.saturating_sub(entry.bytes.len());
1069            }
1070        }
1071    }
1072
1073    fn invalidate_prefix(&self, prefix: &str) {
1074        let mut inner = self.inner.lock().expect("cold cache mutex poisoned");
1075        let keys = inner
1076            .blocks
1077            .keys()
1078            .filter(|key| key.path.starts_with(prefix))
1079            .cloned()
1080            .collect::<Vec<_>>();
1081        for key in keys {
1082            if let Some(entry) = inner.blocks.remove(&key) {
1083                inner.current_bytes = inner.current_bytes.saturating_sub(entry.bytes.len());
1084            }
1085        }
1086    }
1087
1088    #[cfg(test)]
1089    fn block_count(&self) -> usize {
1090        self.inner
1091            .lock()
1092            .expect("cold cache mutex poisoned")
1093            .blocks
1094            .len()
1095    }
1096
1097    fn touch(inner: &mut ColdReadCacheInner, key: ColdCacheKey) {
1098        let generation = Self::next_generation(inner);
1099        if let Some(entry) = inner.blocks.get_mut(&key) {
1100            entry.generation = generation;
1101        }
1102        inner.lru.push_back((key, generation));
1103        Self::compact_lru_if_needed(inner);
1104    }
1105
1106    fn compact_lru_if_needed(inner: &mut ColdReadCacheInner) {
1107        // `touch` appends a fresh (key, generation) on every hit without removing
1108        // the stale prior entry, and `evict_locked` only reclaims those when the
1109        // cache is over `max_bytes`. With a working set at or below the cap but
1110        // repeated hits, the deque would otherwise grow without bound. Rebuild it
1111        // from the live blocks once it bloats past 2x the live entry count —
1112        // amortized O(1) per touch, since each rebuild shrinks it back to
1113        // `blocks.len()` so the next rebuild is `blocks.len()` touches away.
1114        if inner.lru.len() <= inner.blocks.len() * 2 + 16 {
1115            return;
1116        }
1117        let mut live: Vec<(u64, ColdCacheKey)> = inner
1118            .blocks
1119            .iter()
1120            .map(|(key, entry)| (entry.generation, key.clone()))
1121            .collect();
1122        live.sort_unstable_by_key(|(generation, _)| *generation);
1123        inner.lru = live
1124            .into_iter()
1125            .map(|(generation, key)| (key, generation))
1126            .collect();
1127    }
1128
1129    fn next_generation(inner: &mut ColdReadCacheInner) -> u64 {
1130        inner.generation = inner.generation.wrapping_add(1);
1131        inner.generation
1132    }
1133
1134    fn evict_locked(&self, inner: &mut ColdReadCacheInner) {
1135        while inner.current_bytes > self.config.max_bytes {
1136            let Some((key, generation)) = inner.lru.pop_front() else {
1137                break;
1138            };
1139            let Some(entry) = inner.blocks.get(&key) else {
1140                continue;
1141            };
1142            if entry.generation != generation {
1143                continue;
1144            }
1145            let entry = inner
1146                .blocks
1147                .remove(&key)
1148                .expect("cache entry exists after lookup");
1149            inner.current_bytes = inner.current_bytes.saturating_sub(entry.bytes.len());
1150        }
1151    }
1152}
1153
1154fn cold_store_io_error(path: &str, err: opendal::Error) -> io::Error {
1155    io::Error::other(format!("cold object '{path}': {err}"))
1156}
1157
1158#[cfg(not(madsim))]
1159fn cold_object_unix_nanos() -> u128 {
1160    SystemTime::now()
1161        .duration_since(UNIX_EPOCH)
1162        .map(|duration| duration.as_nanos())
1163        .unwrap_or(0)
1164}
1165
1166#[cfg(madsim)]
1167fn cold_object_unix_nanos() -> u128 {
1168    0
1169}
1170
1171pub fn new_cold_chunk_path(
1172    stream_id: &BucketStreamId,
1173    start_offset: u64,
1174    end_offset: u64,
1175) -> String {
1176    let unix_nanos = cold_object_unix_nanos();
1177    let sequence = COLD_CHUNK_SEQUENCE.fetch_add(1, Ordering::Relaxed);
1178    format!(
1179        "{stream_id}/chunks/{start_offset:016x}-{end_offset:016x}-{unix_nanos:032x}-{sequence:016x}.bin"
1180    )
1181}
1182
1183/// The prefix under which all of a stream's cold chunks live. Cold objects are
1184/// stream-exclusive, so removing this prefix reclaims every chunk for a fully
1185/// deleted stream in one sweep. Mirrors the layout of [`new_cold_chunk_path`].
1186pub fn cold_chunk_prefix(stream_id: &BucketStreamId) -> String {
1187    format!("{stream_id}/chunks/")
1188}
1189
1190pub fn new_external_payload_path(stream_id: &BucketStreamId) -> String {
1191    let unix_nanos = cold_object_unix_nanos();
1192    let sequence = COLD_CHUNK_SEQUENCE.fetch_add(1, Ordering::Relaxed);
1193    format!("{stream_id}/external/{unix_nanos:032x}-{sequence:016x}.bin")
1194}
1195
1196/// Reset the global cold-object sequence counter. Only available under
1197/// `cfg(madsim)` so the simulator can clear state between scenarios when
1198/// running multiple seeds in one process (e.g. `Runtime::check_determinism`).
1199#[cfg(madsim)]
1200#[allow(dead_code)]
1201pub fn reset_cold_chunk_sequence_for_sim() {
1202    COLD_CHUNK_SEQUENCE.store(0, Ordering::Relaxed);
1203}
1204
1205#[cfg(test)]
1206mod tests {
1207    use bytes::Bytes;
1208    use ursula_config::config::ColdBackend;
1209
1210    use super::ColdReadCache;
1211    use super::ColdStore;
1212    use crate::ColdConfig;
1213    use crate::ColdReadCacheParams;
1214
1215    fn read_cache_params(store: &ColdStore) -> ColdReadCacheParams {
1216        store
1217            .read_cache
1218            .as_ref()
1219            .map(|cache| cache.config)
1220            .expect("read cache")
1221    }
1222
1223    #[test]
1224    fn try_new_omitted_cache_installs_default_cache() {
1225        let config = ColdConfig {
1226            backend: ColdBackend::Memory,
1227            cache: None,
1228            ..Default::default()
1229        };
1230
1231        let store = ColdStore::try_new(&config).expect("memory cold store");
1232        let cache = read_cache_params(&store);
1233
1234        assert_eq!(cache.max_bytes, 256 * 1024 * 1024);
1235        assert_eq!(cache.block_bytes, 1024 * 1024);
1236        assert_eq!(cache.max_readahead_blocks, 4);
1237    }
1238
1239    #[test]
1240    fn try_new_zero_cache_disables_cache() {
1241        let config = ColdConfig {
1242            backend: ColdBackend::Memory,
1243            cache: Some(ursula_config::ColdCacheConfig {
1244                max_size: ursula_config::HumanSize::bytes(0),
1245                ..Default::default()
1246            }),
1247            ..Default::default()
1248        };
1249
1250        let store = ColdStore::try_new(&config).expect("memory cold store");
1251
1252        assert!(store.read_cache.is_none());
1253    }
1254
1255    #[test]
1256    fn try_new_custom_cache_installs_cache() {
1257        let config = ColdConfig {
1258            backend: ColdBackend::Memory,
1259            cache: Some(ursula_config::ColdCacheConfig {
1260                max_size: ursula_config::HumanSize::mib(7),
1261                block_size: ursula_config::HumanSize::kib(512),
1262                readahead_blocks: 3,
1263            }),
1264            ..Default::default()
1265        };
1266
1267        let store = ColdStore::try_new(&config).expect("memory cold store");
1268        let cache = read_cache_params(&store);
1269
1270        assert_eq!(cache.max_bytes, 7 * 1024 * 1024);
1271        assert_eq!(cache.block_bytes, 512 * 1024);
1272        assert_eq!(cache.max_readahead_blocks, 3);
1273    }
1274
1275    #[test]
1276    fn lru_queue_stays_bounded_under_repeated_hits() {
1277        // Working set fits entirely in the cache (4 blocks == max_bytes), so there
1278        // is no eviction pressure and the recency deque is the only thing that
1279        // could grow. Before compaction it grew by one entry per hit (~40k here);
1280        // it must stay bounded to the live set instead.
1281        let cache = ColdReadCache::new(ColdReadCacheParams {
1282            max_bytes: 4 * 1024,
1283            block_bytes: 1024,
1284            max_readahead_blocks: 0,
1285        });
1286        for index in 0..4 {
1287            cache.insert("p".to_owned(), index, Bytes::from(vec![0u8; 1024]));
1288        }
1289        for _ in 0..10_000 {
1290            for index in 0..4 {
1291                assert!(cache.get("p", index).is_some());
1292            }
1293        }
1294        let inner = cache.inner.lock().expect("cache mutex");
1295        assert_eq!(inner.blocks.len(), 4, "live blocks unchanged");
1296        assert!(
1297            inner.lru.len() <= inner.blocks.len() * 2 + 16,
1298            "lru deque grew unbounded: {} entries for {} live blocks",
1299            inner.lru.len(),
1300            inner.blocks.len(),
1301        );
1302    }
1303
1304    #[test]
1305    fn s3_without_bucket_fails() {
1306        let config = ColdConfig {
1307            backend: ColdBackend::S3,
1308            ..Default::default()
1309        };
1310        let err = ColdStore::try_new(&config).expect_err("s3 without bucket should fail");
1311        assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
1312    }
1313}
ursula_runtime/cold_store.rs

ursula_runtime/
cold_store.rs
