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 futures_util::TryStreamExt;
20use opendal::EntryMode;
21use opendal::Operator;
22use opendal::Scheme;
23use opendal::layers::RetryLayer;
24use opendal::layers::TimeoutLayer;
25use ursula_config::config::ColdBackend;
26use ursula_shard::BucketStreamId;
27use ursula_stream::ColdChunkRef;
28use ursula_stream::ObjectPayloadRef;
29
30use crate::ColdConfig;
31use crate::ColdIndexPageKey;
32
33pub(crate) const DEFAULT_CONTENT_TYPE: &str = "application/octet-stream";
34static COLD_CHUNK_SEQUENCE: CachePadded<AtomicU64> = CachePadded::new(AtomicU64::new(0));
39
40pub(crate) fn apply_s3_encryption(
44 builder: opendal::services::S3,
45 s3: &ursula_config::S3Config,
46) -> io::Result<(opendal::services::S3, &'static str)> {
47 use ursula_config::S3ServerSideEncryption;
48 match s3.server_side_encryption {
49 S3ServerSideEncryption::Aes256 => {
50 if s3.kms_key_id.is_some() {
51 return Err(io::Error::new(
52 io::ErrorKind::InvalidInput,
53 "s3.kms_key_id requires server_side_encryption = \"aws-kms\"",
54 ));
55 }
56 Ok((builder.server_side_encryption_with_s3_key(), "aes256"))
57 }
58 S3ServerSideEncryption::AwsKms => match s3.kms_key_id.as_deref() {
59 Some(key_id) if !key_id.trim().is_empty() => Ok((
60 builder.server_side_encryption_with_customer_managed_kms_key(key_id),
61 "aws-kms",
62 )),
63 _ => Ok((
64 builder.server_side_encryption_with_aws_managed_kms_key(),
65 "aws-kms",
66 )),
67 },
68 S3ServerSideEncryption::None => {
69 if s3.kms_key_id.is_some() {
70 return Err(io::Error::new(
71 io::ErrorKind::InvalidInput,
72 "s3.kms_key_id requires server_side_encryption = \"aws-kms\"",
73 ));
74 }
75 Ok((builder, "none"))
76 }
77 }
78}
79
80pub(crate) fn with_s3_resilience(
96 operator: Operator,
97 timeout: Duration,
98 max_retries: usize,
99) -> Operator {
100 operator
101 .layer(
102 TimeoutLayer::new()
103 .with_timeout(timeout)
104 .with_io_timeout(timeout),
105 )
106 .layer(RetryLayer::new().with_max_times(max_retries).with_jitter())
107}
108
109#[derive(Clone)]
110pub struct ColdStore {
111 info: ColdStoreInfo,
112 operator: Operator,
113 read_cache: Option<Arc<ColdReadCache>>,
114 observer: Arc<Mutex<Option<ColdStoreObserver>>>,
115 fault_policy: Arc<Mutex<Option<ColdStoreFaultPolicy>>>,
116 delay_fn: Arc<Mutex<ColdStoreDelayFn>>,
117}
118
119pub type ColdStoreHandle = Arc<ColdStore>;
120
121type ColdStoreObserver = Arc<dyn Fn(ColdStoreEvent) + Send + Sync>;
122type ColdStoreFaultPolicy =
123 Arc<dyn Fn(&ColdStoreFaultContext) -> Option<ColdStoreFaultEffect> + Send + Sync>;
124type ColdStoreDelayFuture = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
125type ColdStoreDelayFn = Arc<dyn Fn(Duration) -> ColdStoreDelayFuture + Send + Sync>;
126
127#[derive(Debug, Clone, Copy, PartialEq, Eq)]
128pub enum ColdStoreOperation {
129 WriteChunk,
130 DeleteChunk,
131 RemoveAll,
132 ReadObjectRange,
133}
134
135#[derive(Debug, Clone, PartialEq, Eq)]
136pub struct ColdStoreFaultContext {
137 pub operation: ColdStoreOperation,
138 pub stream_id: Option<BucketStreamId>,
139 pub path: String,
140 pub payload_len: Option<usize>,
141 pub read_start_offset: Option<u64>,
142 pub len: Option<usize>,
143 pub object_start: Option<u64>,
144 pub object_end: Option<u64>,
145 pub cached: Option<bool>,
146}
147
148#[derive(Debug, Clone, PartialEq, Eq)]
149pub struct ColdStoreFault {
150 pub message: String,
151}
152
153impl ColdStoreFault {
154 pub fn new(message: impl Into<String>) -> Self {
155 Self {
156 message: message.into(),
157 }
158 }
159}
160
161#[derive(Debug, Clone, Default, PartialEq, Eq)]
162pub struct ColdStoreFaultEffect {
163 pub delay: Option<Duration>,
164 pub error: Option<ColdStoreFault>,
165 pub truncate_read_to: Option<usize>,
166}
167
168impl ColdStoreFaultEffect {
169 pub fn delay(duration: Duration) -> Self {
170 Self {
171 delay: Some(duration),
172 error: None,
173 truncate_read_to: None,
174 }
175 }
176
177 pub fn fail(message: impl Into<String>) -> Self {
178 Self {
179 delay: None,
180 error: Some(ColdStoreFault::new(message)),
181 truncate_read_to: None,
182 }
183 }
184
185 pub fn delay_then_fail(duration: Duration, message: impl Into<String>) -> Self {
186 Self {
187 delay: Some(duration),
188 error: Some(ColdStoreFault::new(message)),
189 truncate_read_to: None,
190 }
191 }
192
193 pub fn truncate_read_to(len: usize) -> Self {
194 Self {
195 delay: None,
196 error: None,
197 truncate_read_to: Some(len),
198 }
199 }
200}
201
202#[derive(Debug, Clone, PartialEq, Eq)]
203pub enum ColdStoreEvent {
204 WriteChunkBegin {
205 path: String,
206 payload_len: usize,
207 },
208 WriteChunkComplete {
209 path: String,
210 object_size: u64,
211 },
212 DeleteChunkBegin {
213 path: String,
214 },
215 DeleteChunkComplete {
216 path: String,
217 },
218 RemoveAllBegin {
219 path: String,
220 },
221 RemoveAllComplete {
222 path: String,
223 },
224 ReadObjectRangeBegin {
225 stream_id: Option<BucketStreamId>,
226 path: String,
227 read_start_offset: u64,
228 len: usize,
229 object_start: u64,
230 object_end: u64,
231 cached: bool,
232 },
233 ReadObjectRangeComplete {
234 stream_id: Option<BucketStreamId>,
235 path: String,
236 read_start_offset: u64,
237 len: usize,
238 returned_len: usize,
239 cached: bool,
240 },
241 FaultInjected {
242 operation: ColdStoreOperation,
243 stream_id: Option<BucketStreamId>,
244 path: String,
245 message: String,
246 },
247 DelayInjected {
248 operation: ColdStoreOperation,
249 stream_id: Option<BucketStreamId>,
250 path: String,
251 delay_ms: u64,
252 },
253 TruncateInjected {
254 stream_id: Option<BucketStreamId>,
255 path: String,
256 requested_len: usize,
257 returned_len: usize,
258 },
259}
260
261impl fmt::Debug for ColdStore {
262 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
263 f.debug_struct("ColdStore")
264 .field("info", &self.info)
265 .field("operator", &self.operator)
266 .field("read_cache", &self.read_cache)
267 .finish_non_exhaustive()
268 }
269}
270
271#[derive(Clone, Debug, PartialEq, Eq)]
272pub struct ColdStoreInfo {
273 pub backend: &'static str,
274 pub root: Option<String>,
275 pub bucket: Option<String>,
276 pub region: Option<String>,
277 pub endpoint: Option<String>,
278 pub encryption: Option<&'static str>,
280}
281
282impl ColdStore {
283 pub async fn list_cold_index_pages(&self) -> io::Result<Vec<ColdIndexPageKey>> {
287 let mut lister = self
288 .operator
289 .lister_with("")
290 .recursive(true)
291 .await
292 .map_err(|err| cold_store_io_error("", err))?;
293 let mut pages = Vec::new();
294 while let Some(entry) = lister
295 .try_next()
296 .await
297 .map_err(|err| cold_store_io_error("", err))?
298 {
299 if entry.metadata().mode() != EntryMode::FILE {
300 continue;
301 }
302 if let Some(key) = parse_cold_index_page_path(entry.path()) {
303 pages.push(key);
304 }
305 }
306 pages.sort_by(|left, right| {
307 left.stream_id
308 .bucket_id
309 .cmp(&right.stream_id.bucket_id)
310 .then_with(|| {
311 left.stream_id
312 .affinity_key
313 .cmp(&right.stream_id.affinity_key)
314 })
315 .then_with(|| left.stream_id.stream_id.cmp(&right.stream_id.stream_id))
316 .then_with(|| left.generation.cmp(&right.generation))
317 .then_with(|| left.page_id.cmp(&right.page_id))
318 });
319 pages.dedup();
320 Ok(pages)
321 }
322
323 pub fn memory() -> io::Result<Self> {
324 let operator = Operator::via_iter(Scheme::Memory, [])
325 .map_err(|err| io::Error::other(err.to_string()))?;
326 Ok(Self::from_operator(operator, ColdStoreInfo {
327 backend: "memory",
328 root: None,
329 bucket: None,
330 region: None,
331 endpoint: None,
332 encryption: None,
333 }))
334 }
335
336 pub fn try_new(config: &ColdConfig) -> io::Result<Self> {
344 let mut store = match config.backend {
345 ColdBackend::None => {
346 return Err(io::Error::new(
347 io::ErrorKind::InvalidInput,
348 "ColdStore::try_new called with backend=none; \
349 the caller should skip construction when cold storage is disabled",
350 ));
351 }
352 ColdBackend::Memory => Self::memory()?,
353 ColdBackend::S3 => Self::s3_from_config(config)?,
354 };
355 let cache = config.cache.clone().unwrap_or_default();
356 if cache.max_size.as_bytes() > 0 {
357 let cache_params = ColdReadCacheParams {
358 max_bytes: cache.max_size.as_bytes() as usize,
359 block_bytes: cache.block_size.as_bytes() as usize,
360 max_readahead_blocks: cache.readahead_blocks,
361 };
362 store = store.with_read_cache(cache_params);
363 }
364 Ok(store)
365 }
366
367 fn s3_from_config(config: &ColdConfig) -> io::Result<Self> {
368 let s3 = config.s3.as_ref().ok_or_else(|| {
369 io::Error::new(
370 io::ErrorKind::InvalidInput,
371 "s3 configuration is required when cold backend is s3",
372 )
373 })?;
374 let bucket = s3.bucket.as_deref().ok_or_else(|| {
375 io::Error::new(
376 io::ErrorKind::InvalidInput,
377 "s3 bucket is required when cold backend is s3",
378 )
379 })?;
380 if bucket.trim().is_empty() {
381 return Err(io::Error::new(
382 io::ErrorKind::InvalidInput,
383 "s3 bucket must not be empty",
384 ));
385 }
386
387 let mut builder = opendal::services::S3::default().bucket(bucket);
388 let mut configured_root = None;
389 if let Some(root) = config.root.as_deref()
390 && !root.trim().is_empty()
391 {
392 builder = builder.root(root);
393 configured_root = Some(root.to_owned());
394 }
395 let mut configured_region = None;
396 if let Some(region) = s3.region.as_deref()
397 && !region.trim().is_empty()
398 {
399 builder = builder.region(region);
400 configured_region = Some(region.to_owned());
401 }
402 let mut configured_endpoint = None;
403 if let Some(endpoint) = s3.endpoint.as_deref()
404 && !endpoint.trim().is_empty()
405 {
406 builder = builder.endpoint(endpoint);
407 configured_endpoint = Some(endpoint.to_owned());
408 }
409 if let Some(access_key_id) = s3.access_key_id.as_deref()
410 && !access_key_id.trim().is_empty()
411 {
412 builder = builder.access_key_id(access_key_id);
413 }
414 if let Some(secret_access_key) = s3.secret_access_key.as_deref()
415 && !secret_access_key.trim().is_empty()
416 {
417 builder = builder.secret_access_key(secret_access_key);
418 }
419 if let Some(session_token) = s3.session_token.as_deref()
420 && !session_token.trim().is_empty()
421 {
422 builder = builder.session_token(session_token);
423 }
424 let (builder, encryption) = apply_s3_encryption(builder, s3)?;
425
426 Ok(Self::from_operator(
427 with_s3_resilience(
428 Operator::new(builder)
429 .map_err(|err| io::Error::other(err.to_string()))?
430 .finish(),
431 s3.timeout.as_duration(),
432 s3.max_retries,
433 ),
434 ColdStoreInfo {
435 backend: "s3",
436 root: configured_root,
437 bucket: Some(bucket.to_owned()),
438 region: configured_region,
439 endpoint: configured_endpoint,
440 encryption: Some(encryption),
441 },
442 ))
443 }
444
445 fn from_operator(operator: Operator, info: ColdStoreInfo) -> Self {
446 Self {
447 info,
448 operator,
449 read_cache: None,
450 observer: Arc::new(Mutex::new(None)),
451 fault_policy: Arc::new(Mutex::new(None)),
452 delay_fn: Arc::new(Mutex::new(default_cold_store_delay_fn())),
453 }
454 }
455
456 pub fn info(&self) -> &ColdStoreInfo {
457 &self.info
458 }
459
460 pub fn with_read_cache(mut self, config: ColdReadCacheParams) -> Self {
461 self.read_cache = Some(Arc::new(ColdReadCache::new(config)));
462 self
463 }
464
465 pub fn without_read_cache(mut self) -> Self {
466 self.read_cache = None;
467 self
468 }
469
470 pub fn set_observer(&self, observer: impl Fn(ColdStoreEvent) + Send + Sync + 'static) {
471 *self.observer.lock().expect("cold store observer mutex") = Some(Arc::new(observer));
472 }
473
474 pub fn set_fault_policy(
475 &self,
476 policy: impl Fn(&ColdStoreFaultContext) -> Option<ColdStoreFaultEffect> + Send + Sync + 'static,
477 ) {
478 *self
479 .fault_policy
480 .lock()
481 .expect("cold store fault policy mutex") = Some(Arc::new(policy));
482 }
483
484 pub fn clear_fault_policy(&self) {
485 *self
486 .fault_policy
487 .lock()
488 .expect("cold store fault policy mutex") = None;
489 }
490
491 pub fn set_delay_fn<F, Fut>(&self, delay_fn: F)
492 where
493 F: Fn(Duration) -> Fut + Send + Sync + 'static,
494 Fut: Future<Output = ()> + Send + 'static,
495 {
496 *self.delay_fn.lock().expect("cold store delay fn mutex") =
497 Arc::new(move |duration| Box::pin(delay_fn(duration)));
498 }
499
500 #[cfg(test)]
501 pub(crate) fn cached_block_count(&self) -> usize {
502 self.read_cache
503 .as_ref()
504 .map(|cache| cache.block_count())
505 .unwrap_or(0)
506 }
507
508 pub async fn write_chunk(&self, path: &str, payload: &[u8]) -> io::Result<u64> {
509 if path.trim().is_empty() {
510 return Err(io::Error::new(
511 io::ErrorKind::InvalidInput,
512 "cold chunk path must not be empty",
513 ));
514 }
515 self.notify(ColdStoreEvent::WriteChunkBegin {
516 path: path.to_owned(),
517 payload_len: payload.len(),
518 });
519 let _applied_fault = self
520 .maybe_apply_fault_effect(ColdStoreFaultContext {
521 operation: ColdStoreOperation::WriteChunk,
522 stream_id: None,
523 path: path.to_owned(),
524 payload_len: Some(payload.len()),
525 read_start_offset: None,
526 len: None,
527 object_start: None,
528 object_end: None,
529 cached: None,
530 })
531 .await?;
532 self.operator
533 .write(path, payload.to_vec())
534 .await
535 .map_err(|err| cold_store_io_error(path, err))?;
536 let object_size = u64::try_from(payload.len()).expect("payload len fits u64");
537 self.notify(ColdStoreEvent::WriteChunkComplete {
538 path: path.to_owned(),
539 object_size,
540 });
541 Ok(object_size)
542 }
543
544 pub(crate) async fn write_cold_index_page(
545 &self,
546 path: &str,
547 payload: &[u8],
548 ) -> io::Result<u64> {
549 if path.trim().is_empty() {
550 return Err(io::Error::new(
551 io::ErrorKind::InvalidInput,
552 "cold index page path must not be empty",
553 ));
554 }
555 self.operator
556 .write(path, payload.to_vec())
557 .await
558 .map_err(|err| cold_store_io_error(path, err))?;
559 Ok(u64::try_from(payload.len()).expect("payload len fits u64"))
560 }
561
562 #[tracing::instrument(name = "cold.read_index", level = "debug", skip_all)]
563 pub(crate) async fn read_cold_index_page(&self, path: &str) -> io::Result<Option<Vec<u8>>> {
564 if path.trim().is_empty() {
565 return Err(io::Error::new(
566 io::ErrorKind::InvalidInput,
567 "cold index page path must not be empty",
568 ));
569 }
570 match self.operator.read(path).await {
571 Ok(bytes) => Ok(Some(bytes.to_bytes().to_vec())),
572 Err(err) if err.kind() == opendal::ErrorKind::NotFound => Ok(None),
573 Err(err) => Err(cold_store_io_error(path, err)),
574 }
575 }
576
577 pub async fn delete_chunk(&self, path: &str) -> io::Result<()> {
578 if path.trim().is_empty() {
579 return Err(io::Error::new(
580 io::ErrorKind::InvalidInput,
581 "cold chunk path must not be empty",
582 ));
583 }
584 self.notify(ColdStoreEvent::DeleteChunkBegin {
585 path: path.to_owned(),
586 });
587 let _applied_fault = self
588 .maybe_apply_fault_effect(ColdStoreFaultContext {
589 operation: ColdStoreOperation::DeleteChunk,
590 stream_id: None,
591 path: path.to_owned(),
592 payload_len: None,
593 read_start_offset: None,
594 len: None,
595 object_start: None,
596 object_end: None,
597 cached: None,
598 })
599 .await?;
600 self.operator
601 .delete(path)
602 .await
603 .map_err(|err| cold_store_io_error(path, err))?;
604 if let Some(cache) = &self.read_cache {
605 cache.invalidate_path(path);
606 }
607 self.notify(ColdStoreEvent::DeleteChunkComplete {
608 path: path.to_owned(),
609 });
610 Ok(())
611 }
612
613 pub async fn remove_all(&self, path: &str) -> io::Result<()> {
614 self.notify(ColdStoreEvent::RemoveAllBegin {
615 path: path.to_owned(),
616 });
617 let _applied_fault = self
618 .maybe_apply_fault_effect(ColdStoreFaultContext {
619 operation: ColdStoreOperation::RemoveAll,
620 stream_id: None,
621 path: path.to_owned(),
622 payload_len: None,
623 read_start_offset: None,
624 len: None,
625 object_start: None,
626 object_end: None,
627 cached: None,
628 })
629 .await?;
630 self.operator
631 .remove_all(path)
632 .await
633 .map_err(|err| cold_store_io_error(path, err))?;
634 if let Some(cache) = &self.read_cache {
635 cache.invalidate_prefix(path);
636 }
637 self.notify(ColdStoreEvent::RemoveAllComplete {
638 path: path.to_owned(),
639 });
640 Ok(())
641 }
642
643 pub async fn read_chunk_range(
644 &self,
645 chunk: &ColdChunkRef,
646 read_start_offset: u64,
647 len: usize,
648 ) -> io::Result<Vec<u8>> {
649 let object = ObjectPayloadRef {
650 start_offset: chunk.start_offset,
651 end_offset: chunk.end_offset,
652 s3_path: chunk.s3_path.clone(),
653 object_size: chunk.object_size,
654 object_offset: chunk.object_offset,
655 };
656 self.read_object_range(&object, read_start_offset, len)
657 .await
658 }
659
660 pub async fn read_object_range_for_stream(
661 &self,
662 stream_id: &BucketStreamId,
663 object: &ObjectPayloadRef,
664 read_start_offset: u64,
665 len: usize,
666 ) -> io::Result<Vec<u8>> {
667 self.read_object_range_inner(Some(stream_id), object, read_start_offset, len)
668 .await
669 }
670
671 pub async fn read_object_range(
672 &self,
673 object: &ObjectPayloadRef,
674 read_start_offset: u64,
675 len: usize,
676 ) -> io::Result<Vec<u8>> {
677 self.read_object_range_inner(None, object, read_start_offset, len)
678 .await
679 }
680
681 #[tracing::instrument(
682 name = "cold.read_chunk",
683 level = "debug",
684 skip_all,
685 fields(
686 start_offset = object.start_offset,
687 end_offset = object.end_offset,
688 object_size = object.object_size,
689 len = len,
690 ),
691 )]
692 async fn read_object_range_inner(
693 &self,
694 stream_id: Option<&BucketStreamId>,
695 object: &ObjectPayloadRef,
696 read_start_offset: u64,
697 len: usize,
698 ) -> io::Result<Vec<u8>> {
699 if len == 0 {
700 return Ok(Vec::new());
701 }
702 let len_u64 = u64::try_from(len).map_err(|_| {
703 io::Error::new(io::ErrorKind::InvalidInput, "cold read length exceeds u64")
704 })?;
705 let read_end = read_start_offset.checked_add(len_u64).ok_or_else(|| {
706 io::Error::new(io::ErrorKind::InvalidInput, "cold read range overflow")
707 })?;
708 if read_start_offset < object.start_offset || read_end > object.end_offset {
709 return Err(io::Error::new(
710 io::ErrorKind::InvalidInput,
711 format!(
712 "cold read range [{read_start_offset}..{read_end}) is outside object segment [{}..{})",
713 object.start_offset, object.end_offset
714 ),
715 ));
716 }
717 let object_start = object
718 .object_offset
719 .checked_add(read_start_offset - object.start_offset)
720 .ok_or_else(|| {
721 io::Error::new(io::ErrorKind::InvalidInput, "cold read range overflow")
722 })?;
723 let object_end = object_start.checked_add(len_u64).ok_or_else(|| {
724 io::Error::new(io::ErrorKind::InvalidInput, "cold read range overflow")
725 })?;
726 if object_end > object.object_size {
727 return Err(io::Error::new(
728 io::ErrorKind::InvalidData,
729 format!(
730 "cold read range [{object_start}..{object_end}) is outside object '{}' size {}",
731 object.s3_path, object.object_size
732 ),
733 ));
734 }
735 let cached = self.read_cache.is_some();
736 self.notify(ColdStoreEvent::ReadObjectRangeBegin {
737 stream_id: stream_id.cloned(),
738 path: object.s3_path.clone(),
739 read_start_offset,
740 len,
741 object_start,
742 object_end,
743 cached,
744 });
745 let applied_fault = self
746 .maybe_apply_fault_effect(ColdStoreFaultContext {
747 operation: ColdStoreOperation::ReadObjectRange,
748 stream_id: stream_id.cloned(),
749 path: object.s3_path.clone(),
750 payload_len: None,
751 read_start_offset: Some(read_start_offset),
752 len: Some(len),
753 object_start: Some(object_start),
754 object_end: Some(object_end),
755 cached: Some(cached),
756 })
757 .await?;
758 let mut bytes = if let Some(cache) = &self.read_cache {
759 let bytes = self
760 .read_object_range_cached(cache, object, object_start, object_end, len)
761 .await?;
762 if let Some(stream_id) = stream_id {
763 let readahead_blocks = cache.record_stream_read(stream_id, read_start_offset, len);
764 if readahead_blocks > 0 {
765 self.spawn_readahead(object.clone(), object_end, readahead_blocks);
766 }
767 }
768 bytes
769 } else {
770 self.read_object_range_uncached(object, object_start, object_end, len)
771 .await?
772 };
773 if let Some(returned_len) = applied_fault.truncate_read_to {
774 let returned_len = returned_len.min(bytes.len());
775 bytes.truncate(returned_len);
776 self.notify(ColdStoreEvent::TruncateInjected {
777 stream_id: stream_id.cloned(),
778 path: object.s3_path.clone(),
779 requested_len: len,
780 returned_len,
781 });
782 }
783 if bytes.len() != len {
784 return Err(io::Error::new(
785 io::ErrorKind::InvalidData,
786 format!(
787 "cold object '{}' returned {} bytes for requested range [{}..{})",
788 object.s3_path,
789 bytes.len(),
790 object_start,
791 object_end
792 ),
793 ));
794 }
795 self.notify(ColdStoreEvent::ReadObjectRangeComplete {
796 stream_id: stream_id.cloned(),
797 path: object.s3_path.clone(),
798 read_start_offset,
799 len,
800 returned_len: bytes.len(),
801 cached,
802 });
803 Ok(bytes)
804 }
805
806 async fn read_object_range_uncached(
807 &self,
808 object: &ObjectPayloadRef,
809 object_start: u64,
810 object_end: u64,
811 len: usize,
812 ) -> io::Result<Vec<u8>> {
813 let bytes = self
814 .operator
815 .read_with(&object.s3_path)
816 .range(object_start..object_end)
817 .await
818 .map_err(|err| cold_store_io_error(&object.s3_path, err))?
819 .to_bytes();
820 if bytes.len() != len {
821 return Err(io::Error::new(
822 io::ErrorKind::InvalidData,
823 format!(
824 "cold object '{}' returned {} bytes for requested range [{}..{})",
825 object.s3_path,
826 bytes.len(),
827 object_start,
828 object_end
829 ),
830 ));
831 }
832 Ok(bytes.to_vec())
833 }
834
835 async fn read_object_range_cached(
836 &self,
837 cache: &ColdReadCache,
838 object: &ObjectPayloadRef,
839 object_start: u64,
840 object_end: u64,
841 len: usize,
842 ) -> io::Result<Vec<u8>> {
843 let mut payload = Vec::with_capacity(len);
844 let block_size = cache.block_size();
845 let first_block = object_start / block_size;
846 let last_block = (object_end - 1) / block_size;
847 for block_index in first_block..=last_block {
848 let block_start = block_index * block_size;
849 let block_end = block_start
850 .saturating_add(block_size)
851 .min(object.object_size);
852 let block = self
853 .read_cached_block(
854 cache,
855 object.s3_path.clone(),
856 object.object_size,
857 block_index,
858 block_start,
859 block_end,
860 )
861 .await?;
862 let slice_start = usize::try_from(object_start.max(block_start) - block_start)
863 .expect("cache slice start fits usize");
864 let slice_end = usize::try_from(object_end.min(block_end) - block_start)
865 .expect("cache slice end fits usize");
866 payload.extend_from_slice(&block.slice(slice_start..slice_end));
867 }
868 if payload.len() != len {
869 return Err(io::Error::new(
870 io::ErrorKind::InvalidData,
871 format!(
872 "cold object '{}' returned {} bytes for requested range [{}..{})",
873 object.s3_path,
874 payload.len(),
875 object_start,
876 object_end
877 ),
878 ));
879 }
880 Ok(payload)
881 }
882
883 async fn read_cached_block(
884 &self,
885 cache: &ColdReadCache,
886 path: String,
887 object_size: u64,
888 block_index: u64,
889 block_start: u64,
890 block_end: u64,
891 ) -> io::Result<Bytes> {
892 if let Some(bytes) = cache.get(&path, block_index) {
893 return Ok(bytes);
894 }
895 let bytes = self
896 .operator
897 .read_with(&path)
898 .range(block_start..block_end)
899 .await
900 .map_err(|err| cold_store_io_error(&path, err))?
901 .to_bytes();
902 let expected_len = usize::try_from(block_end - block_start).map_err(|_| {
903 io::Error::new(
904 io::ErrorKind::InvalidData,
905 "cold cache block length exceeds usize",
906 )
907 })?;
908 if bytes.len() != expected_len {
909 return Err(io::Error::new(
910 io::ErrorKind::InvalidData,
911 format!(
912 "cold object '{path}' returned {} bytes for cache block [{}..{}) of object size {object_size}",
913 bytes.len(),
914 block_start,
915 block_end
916 ),
917 ));
918 }
919 cache.insert(path, block_index, bytes.clone());
920 Ok(bytes)
921 }
922
923 fn spawn_readahead(&self, object: ObjectPayloadRef, object_end: u64, readahead_blocks: usize) {
924 let Some(cache) = self.read_cache.clone() else {
925 return;
926 };
927 let block_size = cache.block_size();
928 let mut block_index = object_end.div_ceil(block_size);
929 let store = self.clone();
930 crate::rt::spawn(async move {
931 for _ in 0..readahead_blocks {
932 let block_start = block_index * block_size;
933 if block_start >= object.object_size {
934 break;
935 }
936 let block_end = block_start
937 .saturating_add(block_size)
938 .min(object.object_size);
939 if cache.get(&object.s3_path, block_index).is_none() {
940 let _ = store
941 .read_cached_block(
942 &cache,
943 object.s3_path.clone(),
944 object.object_size,
945 block_index,
946 block_start,
947 block_end,
948 )
949 .await;
950 }
951 block_index += 1;
952 }
953 });
954 }
955
956 fn notify(&self, event: ColdStoreEvent) {
957 let observer = self
958 .observer
959 .lock()
960 .expect("cold store observer mutex")
961 .clone();
962 if let Some(observer) = observer {
963 observer(event);
964 }
965 }
966
967 async fn maybe_apply_fault_effect(
968 &self,
969 context: ColdStoreFaultContext,
970 ) -> io::Result<ColdStoreAppliedFault> {
971 let policy = self
972 .fault_policy
973 .lock()
974 .expect("cold store fault policy mutex")
975 .clone();
976 let Some(policy) = policy else {
977 return Ok(ColdStoreAppliedFault::default());
978 };
979 let Some(effect) = policy(&context) else {
980 return Ok(ColdStoreAppliedFault::default());
981 };
982 if let Some(delay) = effect.delay {
983 self.notify(ColdStoreEvent::DelayInjected {
984 operation: context.operation,
985 stream_id: context.stream_id.clone(),
986 path: context.path.clone(),
987 delay_ms: duration_ms(delay),
988 });
989 let delay_fn = self
990 .delay_fn
991 .lock()
992 .expect("cold store delay fn mutex")
993 .clone();
994 delay_fn(delay).await;
995 }
996 if let Some(fault) = effect.error {
997 self.notify(ColdStoreEvent::FaultInjected {
998 operation: context.operation,
999 stream_id: context.stream_id,
1000 path: context.path.clone(),
1001 message: fault.message.clone(),
1002 });
1003 return Err(io::Error::other(format!(
1004 "cold store fault injected for {} '{}': {}",
1005 context.operation.as_str(),
1006 context.path,
1007 fault.message
1008 )));
1009 }
1010 Ok(ColdStoreAppliedFault {
1011 truncate_read_to: effect.truncate_read_to,
1012 })
1013 }
1014}
1015
1016fn parse_cold_index_page_path(path: &str) -> Option<ColdIndexPageKey> {
1017 let parts = path.split('/').collect::<Vec<_>>();
1018 let (stream_id, generation, page_id) = match parts.as_slice() {
1019 [bucket, stream, "cold-index", generation, page] => {
1020 (BucketStreamId::new(*bucket, *stream), *generation, *page)
1021 }
1022 [bucket, affinity, stream, "cold-index", generation, page] => (
1023 BucketStreamId::with_affinity(*bucket, *affinity, *stream),
1024 *generation,
1025 *page,
1026 ),
1027 _ => return None,
1028 };
1029 if stream_id.bucket_id.is_empty()
1030 || stream_id
1031 .affinity_key
1032 .as_ref()
1033 .is_some_and(String::is_empty)
1034 || stream_id.stream_id.is_empty()
1035 {
1036 return None;
1037 }
1038 Some(ColdIndexPageKey {
1039 stream_id,
1040 generation: generation.parse().ok()?,
1041 page_id: page_id.strip_suffix(".idx")?.parse().ok()?,
1042 })
1043}
1044
1045#[derive(Debug, Default)]
1046struct ColdStoreAppliedFault {
1047 truncate_read_to: Option<usize>,
1048}
1049
1050impl ColdStoreOperation {
1051 fn as_str(self) -> &'static str {
1052 match self {
1053 Self::WriteChunk => "write_chunk",
1054 Self::DeleteChunk => "delete_chunk",
1055 Self::RemoveAll => "remove_all",
1056 Self::ReadObjectRange => "read_object_range",
1057 }
1058 }
1059}
1060
1061fn default_cold_store_delay_fn() -> ColdStoreDelayFn {
1062 Arc::new(|duration| Box::pin(crate::rt::time::sleep(duration)))
1063}
1064
1065fn duration_ms(duration: Duration) -> u64 {
1066 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
1067}
1068
1069#[derive(Debug, Clone, Copy)]
1074pub struct ColdReadCacheParams {
1075 pub max_bytes: usize,
1076 pub block_bytes: usize,
1077 pub max_readahead_blocks: usize,
1078}
1079
1080#[derive(Debug)]
1081struct ColdReadCache {
1082 config: ColdReadCacheParams,
1083 inner: Mutex<ColdReadCacheInner>,
1084}
1085
1086#[derive(Debug, Default)]
1087struct ColdReadCacheInner {
1088 blocks: HashMap<ColdCacheKey, ColdCacheEntry>,
1089 lru: VecDeque<(ColdCacheKey, u64)>,
1090 current_bytes: usize,
1091 generation: u64,
1092 readers: HashMap<BucketStreamId, StreamReadState>,
1093}
1094
1095#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1096struct ColdCacheKey {
1097 path: String,
1098 block_index: u64,
1099}
1100
1101#[derive(Debug)]
1102struct ColdCacheEntry {
1103 bytes: Bytes,
1104 generation: u64,
1105}
1106
1107#[derive(Debug, Default)]
1108struct StreamReadState {
1109 next_offset: u64,
1110 sequential_score: usize,
1111}
1112
1113impl ColdReadCache {
1114 fn new(config: ColdReadCacheParams) -> Self {
1115 let block_bytes = config.block_bytes.max(1);
1116 Self {
1117 config: ColdReadCacheParams {
1118 max_bytes: config.max_bytes,
1119 block_bytes,
1120 max_readahead_blocks: config.max_readahead_blocks,
1121 },
1122 inner: Mutex::new(ColdReadCacheInner::default()),
1123 }
1124 }
1125
1126 fn block_size(&self) -> u64 {
1127 u64::try_from(self.config.block_bytes.max(1)).expect("cache block size fits u64")
1128 }
1129
1130 fn get(&self, path: &str, block_index: u64) -> Option<Bytes> {
1131 let mut inner = self.inner.lock().expect("cold cache mutex poisoned");
1132 let key = ColdCacheKey {
1133 path: path.to_owned(),
1134 block_index,
1135 };
1136 let bytes = inner.blocks.get(&key)?.bytes.clone();
1137 Self::touch(&mut inner, key);
1138 Some(bytes)
1139 }
1140
1141 fn insert(&self, path: String, block_index: u64, bytes: Bytes) {
1142 if bytes.len() > self.config.max_bytes || self.config.max_bytes == 0 {
1143 return;
1144 }
1145 let mut inner = self.inner.lock().expect("cold cache mutex poisoned");
1146 let key = ColdCacheKey { path, block_index };
1147 if let Some(previous) = inner.blocks.remove(&key) {
1148 inner.current_bytes = inner.current_bytes.saturating_sub(previous.bytes.len());
1149 }
1150 let generation = Self::next_generation(&mut inner);
1151 inner.current_bytes = inner.current_bytes.saturating_add(bytes.len());
1152 inner
1153 .blocks
1154 .insert(key.clone(), ColdCacheEntry { bytes, generation });
1155 inner.lru.push_back((key, generation));
1156 self.evict_locked(&mut inner);
1157 Self::compact_lru_if_needed(&mut inner);
1158 }
1159
1160 fn record_stream_read(
1161 &self,
1162 stream_id: &BucketStreamId,
1163 read_start_offset: u64,
1164 len: usize,
1165 ) -> usize {
1166 let mut inner = self.inner.lock().expect("cold cache mutex poisoned");
1167 let state = inner.readers.entry(stream_id.clone()).or_default();
1168 if read_start_offset == state.next_offset {
1169 state.sequential_score = state
1170 .sequential_score
1171 .saturating_add(1)
1172 .min(self.config.max_readahead_blocks);
1173 } else {
1174 state.sequential_score = 0;
1175 }
1176 state.next_offset =
1177 read_start_offset.saturating_add(u64::try_from(len).unwrap_or(u64::MAX));
1178 state.sequential_score.min(self.config.max_readahead_blocks)
1179 }
1180
1181 fn invalidate_path(&self, path: &str) {
1182 let mut inner = self.inner.lock().expect("cold cache mutex poisoned");
1183 let keys = inner
1184 .blocks
1185 .keys()
1186 .filter(|key| key.path == path)
1187 .cloned()
1188 .collect::<Vec<_>>();
1189 for key in keys {
1190 if let Some(entry) = inner.blocks.remove(&key) {
1191 inner.current_bytes = inner.current_bytes.saturating_sub(entry.bytes.len());
1192 }
1193 }
1194 }
1195
1196 fn invalidate_prefix(&self, prefix: &str) {
1197 let mut inner = self.inner.lock().expect("cold cache mutex poisoned");
1198 let keys = inner
1199 .blocks
1200 .keys()
1201 .filter(|key| key.path.starts_with(prefix))
1202 .cloned()
1203 .collect::<Vec<_>>();
1204 for key in keys {
1205 if let Some(entry) = inner.blocks.remove(&key) {
1206 inner.current_bytes = inner.current_bytes.saturating_sub(entry.bytes.len());
1207 }
1208 }
1209 }
1210
1211 #[cfg(test)]
1212 fn block_count(&self) -> usize {
1213 self.inner
1214 .lock()
1215 .expect("cold cache mutex poisoned")
1216 .blocks
1217 .len()
1218 }
1219
1220 fn touch(inner: &mut ColdReadCacheInner, key: ColdCacheKey) {
1221 let generation = Self::next_generation(inner);
1222 if let Some(entry) = inner.blocks.get_mut(&key) {
1223 entry.generation = generation;
1224 }
1225 inner.lru.push_back((key, generation));
1226 Self::compact_lru_if_needed(inner);
1227 }
1228
1229 fn compact_lru_if_needed(inner: &mut ColdReadCacheInner) {
1230 if inner.lru.len() <= inner.blocks.len() * 2 + 16 {
1238 return;
1239 }
1240 let mut live: Vec<(u64, ColdCacheKey)> = inner
1241 .blocks
1242 .iter()
1243 .map(|(key, entry)| (entry.generation, key.clone()))
1244 .collect();
1245 live.sort_unstable_by_key(|(generation, _)| *generation);
1246 inner.lru = live
1247 .into_iter()
1248 .map(|(generation, key)| (key, generation))
1249 .collect();
1250 }
1251
1252 fn next_generation(inner: &mut ColdReadCacheInner) -> u64 {
1253 inner.generation = inner.generation.wrapping_add(1);
1254 inner.generation
1255 }
1256
1257 fn evict_locked(&self, inner: &mut ColdReadCacheInner) {
1258 while inner.current_bytes > self.config.max_bytes {
1259 let Some((key, generation)) = inner.lru.pop_front() else {
1260 break;
1261 };
1262 let Some(entry) = inner.blocks.get(&key) else {
1263 continue;
1264 };
1265 if entry.generation != generation {
1266 continue;
1267 }
1268 let entry = inner
1269 .blocks
1270 .remove(&key)
1271 .expect("cache entry exists after lookup");
1272 inner.current_bytes = inner.current_bytes.saturating_sub(entry.bytes.len());
1273 }
1274 }
1275}
1276
1277fn cold_store_io_error(path: &str, err: opendal::Error) -> io::Error {
1278 io::Error::other(format!("cold object '{path}': {err}"))
1279}
1280
1281#[cfg(not(madsim))]
1282fn cold_object_unix_nanos() -> u128 {
1283 SystemTime::now()
1284 .duration_since(UNIX_EPOCH)
1285 .map(|duration| duration.as_nanos())
1286 .unwrap_or(0)
1287}
1288
1289#[cfg(madsim)]
1290fn cold_object_unix_nanos() -> u128 {
1291 0
1292}
1293
1294pub fn new_cold_chunk_path(
1295 stream_id: &BucketStreamId,
1296 start_offset: u64,
1297 end_offset: u64,
1298) -> String {
1299 let unix_nanos = cold_object_unix_nanos();
1300 let sequence = COLD_CHUNK_SEQUENCE.fetch_add(1, Ordering::Relaxed);
1301 format!(
1302 "{stream_id}/chunks/{start_offset:016x}-{end_offset:016x}-{unix_nanos:032x}-{sequence:016x}.bin"
1303 )
1304}
1305
1306pub fn new_cold_pack_path(raft_group_id: u32) -> String {
1307 let unix_nanos = cold_object_unix_nanos();
1308 let sequence = COLD_CHUNK_SEQUENCE.fetch_add(1, Ordering::Relaxed);
1309 format!("_packs/{raft_group_id:08x}/{unix_nanos:032x}-{sequence:016x}.bin")
1310}
1311
1312pub fn cold_chunk_prefix(stream_id: &BucketStreamId) -> String {
1316 format!("{stream_id}/chunks/")
1317}
1318
1319pub fn new_external_payload_path(stream_id: &BucketStreamId) -> String {
1320 let unix_nanos = cold_object_unix_nanos();
1321 let sequence = COLD_CHUNK_SEQUENCE.fetch_add(1, Ordering::Relaxed);
1322 format!("{stream_id}/external/{unix_nanos:032x}-{sequence:016x}.bin")
1323}
1324
1325#[cfg(madsim)]
1329#[allow(dead_code)]
1330pub fn reset_cold_chunk_sequence_for_sim() {
1331 COLD_CHUNK_SEQUENCE.store(0, Ordering::Relaxed);
1332}
1333
1334#[cfg(test)]
1335mod tests {
1336 use bytes::Bytes;
1337 use ursula_config::config::ColdBackend;
1338 use ursula_shard::BucketStreamId;
1339
1340 use super::ColdReadCache;
1341 use super::ColdStore;
1342 use super::parse_cold_index_page_path;
1343 use crate::ColdConfig;
1344 use crate::ColdReadCacheParams;
1345
1346 fn read_cache_params(store: &ColdStore) -> ColdReadCacheParams {
1347 store
1348 .read_cache
1349 .as_ref()
1350 .map(|cache| cache.config)
1351 .expect("read cache")
1352 }
1353
1354 #[test]
1355 fn cold_index_paths_preserve_optional_affinity() {
1356 let plain = parse_cold_index_page_path(
1357 "benchcmp/journal/cold-index/00000000000000000007/00000000000000000042.idx",
1358 )
1359 .expect("plain path");
1360 assert_eq!(plain.stream_id, BucketStreamId::new("benchcmp", "journal"));
1361
1362 let grouped = parse_cold_index_page_path(
1363 "benchcmp/run-42/journal/cold-index/00000000000000000007/00000000000000000042.idx",
1364 )
1365 .expect("grouped path");
1366 assert_eq!(
1367 grouped.stream_id,
1368 BucketStreamId::with_affinity("benchcmp", "run-42", "journal")
1369 );
1370 }
1371
1372 #[test]
1373 fn try_new_omitted_cache_installs_default_cache() {
1374 let config = ColdConfig {
1375 backend: ColdBackend::Memory,
1376 cache: None,
1377 ..Default::default()
1378 };
1379
1380 let store = ColdStore::try_new(&config).expect("memory cold store");
1381 let cache = read_cache_params(&store);
1382
1383 assert_eq!(cache.max_bytes, 256 * 1024 * 1024);
1384 assert_eq!(cache.block_bytes, 1024 * 1024);
1385 assert_eq!(cache.max_readahead_blocks, 4);
1386 }
1387
1388 #[test]
1389 fn try_new_zero_cache_disables_cache() {
1390 let config = ColdConfig {
1391 backend: ColdBackend::Memory,
1392 cache: Some(ursula_config::ColdCacheConfig {
1393 max_size: ursula_config::HumanSize::bytes(0),
1394 ..Default::default()
1395 }),
1396 ..Default::default()
1397 };
1398
1399 let store = ColdStore::try_new(&config).expect("memory cold store");
1400
1401 assert!(store.read_cache.is_none());
1402 }
1403
1404 #[test]
1405 fn try_new_custom_cache_installs_cache() {
1406 let config = ColdConfig {
1407 backend: ColdBackend::Memory,
1408 cache: Some(ursula_config::ColdCacheConfig {
1409 max_size: ursula_config::HumanSize::mib(7),
1410 block_size: ursula_config::HumanSize::kib(512),
1411 readahead_blocks: 3,
1412 }),
1413 ..Default::default()
1414 };
1415
1416 let store = ColdStore::try_new(&config).expect("memory cold store");
1417 let cache = read_cache_params(&store);
1418
1419 assert_eq!(cache.max_bytes, 7 * 1024 * 1024);
1420 assert_eq!(cache.block_bytes, 512 * 1024);
1421 assert_eq!(cache.max_readahead_blocks, 3);
1422 }
1423
1424 fn s3_test_config(
1425 encryption: ursula_config::S3ServerSideEncryption,
1426 kms_key_id: Option<&str>,
1427 ) -> ColdConfig {
1428 ColdConfig {
1429 backend: ColdBackend::S3,
1430 s3: Some(ursula_config::S3Config {
1431 bucket: Some("test-bucket".to_owned()),
1432 region: Some("us-east-1".to_owned()),
1433 server_side_encryption: encryption,
1434 kms_key_id: kms_key_id.map(str::to_owned),
1435 ..Default::default()
1436 }),
1437 ..Default::default()
1438 }
1439 }
1440
1441 #[test]
1442 fn s3_store_defaults_to_sse_s3_and_reports_it() {
1443 let store = ColdStore::try_new(&s3_test_config(
1444 ursula_config::S3ServerSideEncryption::Aes256,
1445 None,
1446 ))
1447 .expect("s3 cold store");
1448 assert_eq!(store.info().encryption, Some("aes256"));
1449 }
1450
1451 #[test]
1452 fn s3_store_reports_kms_and_disabled_modes() {
1453 let kms = ColdStore::try_new(&s3_test_config(
1454 ursula_config::S3ServerSideEncryption::AwsKms,
1455 Some("arn:aws:kms:us-east-1:111122223333:key/test"),
1456 ))
1457 .expect("s3 cold store with kms");
1458 assert_eq!(kms.info().encryption, Some("aws-kms"));
1459
1460 let disabled = ColdStore::try_new(&s3_test_config(
1461 ursula_config::S3ServerSideEncryption::None,
1462 None,
1463 ))
1464 .expect("s3 cold store without sse");
1465 assert_eq!(disabled.info().encryption, Some("none"));
1466 }
1467
1468 #[test]
1469 fn kms_key_without_kms_mode_is_rejected() {
1470 let err = ColdStore::try_new(&s3_test_config(
1471 ursula_config::S3ServerSideEncryption::Aes256,
1472 Some("arn:aws:kms:us-east-1:111122223333:key/test"),
1473 ))
1474 .expect_err("kms key without aws-kms mode");
1475 assert!(err.to_string().contains("aws-kms"), "got: {err}");
1476 }
1477
1478 #[test]
1479 fn lru_queue_stays_bounded_under_repeated_hits() {
1480 let cache = ColdReadCache::new(ColdReadCacheParams {
1485 max_bytes: 4 * 1024,
1486 block_bytes: 1024,
1487 max_readahead_blocks: 0,
1488 });
1489 for index in 0..4 {
1490 cache.insert("p".to_owned(), index, Bytes::from(vec![0u8; 1024]));
1491 }
1492 for _ in 0..10_000 {
1493 for index in 0..4 {
1494 assert!(cache.get("p", index).is_some());
1495 }
1496 }
1497 let inner = cache.inner.lock().expect("cache mutex");
1498 assert_eq!(inner.blocks.len(), 4, "live blocks unchanged");
1499 assert!(
1500 inner.lru.len() <= inner.blocks.len() * 2 + 16,
1501 "lru deque grew unbounded: {} entries for {} live blocks",
1502 inner.lru.len(),
1503 inner.blocks.len(),
1504 );
1505 }
1506
1507 #[test]
1508 fn s3_without_bucket_fails() {
1509 let config = ColdConfig {
1510 backend: ColdBackend::S3,
1511 ..Default::default()
1512 };
1513 let err = ColdStore::try_new(&config).expect_err("s3 without bucket should fail");
1514 assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
1515 }
1516}