ursula_runtime/
cold_index.rs

1use std::collections::HashMap;
2use std::collections::VecDeque;
3use std::future::Future;
4use std::io;
5use std::pin::Pin;
6use std::sync::Arc;
7use std::sync::Mutex;
8
9use ursula_shard::BucketStreamId;
10use ursula_stream::ColdChunkRef;
11use ursula_stream::ExternalPayloadRef;
12use ursula_stream::ObjectPayloadRef;
13use ursula_stream::StreamReadColdIndexSegment;
14
15use crate::cold_store::ColdStoreHandle;
16
17pub type ColdIndexPageStoreFuture<'a, T> = Pin<Box<dyn Future<Output = io::Result<T>> + Send + 'a>>;
18
19const COLD_INDEX_PAGE_MAGIC: &[u8; 8] = b"UCIDX001";
20const COLD_INDEX_PAGE_VERSION: u16 = 1;
21const COLD_INDEX_ENTRY_COLD_CHUNK: u8 = 1;
22const COLD_INDEX_ENTRY_EXTERNAL_SEGMENT: u8 = 2;
23const FNV64_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
24const FNV64_PRIME: u64 = 0x0000_0100_0000_01b3;
25
26#[derive(Debug, Clone, PartialEq, Eq, Hash)]
27pub struct ColdIndexPageKey {
28    pub stream_id: BucketStreamId,
29    pub generation: u64,
30    pub page_id: u64,
31}
32
33impl ColdIndexPageKey {
34    pub fn path(&self) -> String {
35        format!(
36            "{}/{}/cold-index/{:020}/{:020}.idx",
37            self.stream_id.bucket_id, self.stream_id.stream_id, self.generation, self.page_id
38        )
39    }
40}
41
42pub fn cold_index_prefix(stream_id: &BucketStreamId) -> String {
43    format!(
44        "{}/{}/cold-index/",
45        stream_id.bucket_id, stream_id.stream_id
46    )
47}
48
49#[derive(Debug, Clone, PartialEq, Eq)]
50pub struct ColdIndexPage {
51    pub start_offset: u64,
52    pub end_offset: u64,
53    pub cold_chunks: Vec<ColdChunkRef>,
54    pub external_segments: Vec<ObjectPayloadRef>,
55}
56
57impl ColdIndexPage {
58    pub fn covers(&self, offset: u64) -> bool {
59        self.start_offset <= offset && offset < self.end_offset
60    }
61}
62
63#[derive(Debug, Clone)]
64pub struct ColdIndexPageRollback {
65    key: ColdIndexPageKey,
66    previous: Option<ColdIndexPage>,
67    written_chunk: ColdChunkRef,
68}
69
70fn encode_page(key: &ColdIndexPageKey, page: &ColdIndexPage) -> Vec<u8> {
71    let mut body = Vec::new();
72    put_string(&mut body, &key.stream_id.bucket_id);
73    put_string(&mut body, &key.stream_id.stream_id);
74    put_u64(&mut body, key.generation);
75    put_u64(&mut body, key.page_id);
76    put_u64(&mut body, page.start_offset);
77    put_u64(&mut body, page.end_offset);
78    put_u32(
79        &mut body,
80        u32::try_from(page.cold_chunks.len()).expect("cold index cold chunk count fits u32"),
81    );
82    for chunk in &page.cold_chunks {
83        put_u8(&mut body, COLD_INDEX_ENTRY_COLD_CHUNK);
84        put_u64(&mut body, chunk.start_offset);
85        put_u64(&mut body, chunk.end_offset);
86        put_u64(&mut body, chunk.object_size);
87        put_string(&mut body, &chunk.s3_path);
88    }
89    put_u32(
90        &mut body,
91        u32::try_from(page.external_segments.len())
92            .expect("cold index external segment count fits u32"),
93    );
94    for object in &page.external_segments {
95        put_u8(&mut body, COLD_INDEX_ENTRY_EXTERNAL_SEGMENT);
96        put_u64(&mut body, object.start_offset);
97        put_u64(&mut body, object.end_offset);
98        put_u64(&mut body, object.object_size);
99        put_string(&mut body, &object.s3_path);
100    }
101
102    let mut bytes = Vec::with_capacity(COLD_INDEX_PAGE_MAGIC.len() + 2 + 4 + body.len() + 8);
103    bytes.extend_from_slice(COLD_INDEX_PAGE_MAGIC);
104    put_u16(&mut bytes, COLD_INDEX_PAGE_VERSION);
105    put_u32(
106        &mut bytes,
107        u32::try_from(body.len()).expect("cold index page body len fits u32"),
108    );
109    bytes.extend_from_slice(&body);
110    put_u64(&mut bytes, checksum64(&body));
111    bytes
112}
113
114fn decode_page(key: &ColdIndexPageKey, bytes: &[u8]) -> io::Result<ColdIndexPage> {
115    let mut cursor = Cursor::new(bytes);
116    let magic = cursor.read_exact(COLD_INDEX_PAGE_MAGIC.len())?;
117    if magic != COLD_INDEX_PAGE_MAGIC {
118        return Err(io::Error::new(
119            io::ErrorKind::InvalidData,
120            "cold index page has invalid magic",
121        ));
122    }
123    let version = cursor.read_u16()?;
124    if version != COLD_INDEX_PAGE_VERSION {
125        return Err(io::Error::new(
126            io::ErrorKind::InvalidData,
127            format!("unsupported cold index page version {version}"),
128        ));
129    }
130    let body_len = usize::try_from(cursor.read_u32()?).expect("u32 fits usize");
131    let body = cursor.read_exact(body_len)?;
132    let expected_checksum = cursor.read_u64()?;
133    if cursor.remaining() != 0 {
134        return Err(io::Error::new(
135            io::ErrorKind::InvalidData,
136            "cold index page has trailing bytes",
137        ));
138    }
139    let actual_checksum = checksum64(body);
140    if actual_checksum != expected_checksum {
141        return Err(io::Error::new(
142            io::ErrorKind::InvalidData,
143            "cold index page checksum mismatch",
144        ));
145    }
146
147    let mut body = Cursor::new(body);
148    let bucket_id = body.read_string()?;
149    let stream_id = body.read_string()?;
150    let generation = body.read_u64()?;
151    let page_id = body.read_u64()?;
152    if bucket_id != key.stream_id.bucket_id
153        || stream_id != key.stream_id.stream_id
154        || generation != key.generation
155        || page_id != key.page_id
156    {
157        return Err(io::Error::new(
158            io::ErrorKind::InvalidData,
159            "cold index page key metadata mismatch",
160        ));
161    }
162    let start_offset = body.read_u64()?;
163    let end_offset = body.read_u64()?;
164    let cold_chunk_count = body.read_u32()?;
165    let mut cold_chunks =
166        Vec::with_capacity(usize::try_from(cold_chunk_count).expect("u32 fits usize"));
167    for _ in 0..cold_chunk_count {
168        let tag = body.read_u8()?;
169        if tag != COLD_INDEX_ENTRY_COLD_CHUNK {
170            return Err(io::Error::new(
171                io::ErrorKind::InvalidData,
172                "cold index page expected cold chunk entry",
173            ));
174        }
175        cold_chunks.push(ColdChunkRef {
176            start_offset: body.read_u64()?,
177            end_offset: body.read_u64()?,
178            object_size: body.read_u64()?,
179            s3_path: body.read_string()?,
180        });
181    }
182    let external_segment_count = body.read_u32()?;
183    let mut external_segments =
184        Vec::with_capacity(usize::try_from(external_segment_count).expect("u32 fits usize"));
185    for _ in 0..external_segment_count {
186        let tag = body.read_u8()?;
187        if tag != COLD_INDEX_ENTRY_EXTERNAL_SEGMENT {
188            return Err(io::Error::new(
189                io::ErrorKind::InvalidData,
190                "cold index page expected external segment entry",
191            ));
192        }
193        external_segments.push(ObjectPayloadRef {
194            start_offset: body.read_u64()?,
195            end_offset: body.read_u64()?,
196            object_size: body.read_u64()?,
197            s3_path: body.read_string()?,
198        });
199    }
200    if body.remaining() != 0 {
201        return Err(io::Error::new(
202            io::ErrorKind::InvalidData,
203            "cold index page body has trailing bytes",
204        ));
205    }
206    Ok(ColdIndexPage {
207        start_offset,
208        end_offset,
209        cold_chunks,
210        external_segments,
211    })
212}
213
214fn put_u8(out: &mut Vec<u8>, value: u8) {
215    out.push(value);
216}
217
218fn put_u16(out: &mut Vec<u8>, value: u16) {
219    out.extend_from_slice(&value.to_le_bytes());
220}
221
222fn put_u32(out: &mut Vec<u8>, value: u32) {
223    out.extend_from_slice(&value.to_le_bytes());
224}
225
226fn put_u64(out: &mut Vec<u8>, value: u64) {
227    out.extend_from_slice(&value.to_le_bytes());
228}
229
230fn put_string(out: &mut Vec<u8>, value: &str) {
231    put_u32(
232        out,
233        u32::try_from(value.len()).expect("cold index string len fits u32"),
234    );
235    out.extend_from_slice(value.as_bytes());
236}
237
238fn checksum64(bytes: &[u8]) -> u64 {
239    let mut hash = FNV64_OFFSET_BASIS;
240    for byte in bytes {
241        hash ^= u64::from(*byte);
242        hash = hash.wrapping_mul(FNV64_PRIME);
243    }
244    hash
245}
246
247struct Cursor<'a> {
248    bytes: &'a [u8],
249    offset: usize,
250}
251
252impl<'a> Cursor<'a> {
253    fn new(bytes: &'a [u8]) -> Self {
254        Self { bytes, offset: 0 }
255    }
256
257    fn remaining(&self) -> usize {
258        self.bytes.len().saturating_sub(self.offset)
259    }
260
261    fn read_exact(&mut self, len: usize) -> io::Result<&'a [u8]> {
262        let end = self.offset.checked_add(len).ok_or_else(|| {
263            io::Error::new(
264                io::ErrorKind::InvalidData,
265                "cold index page offset overflow",
266            )
267        })?;
268        if end > self.bytes.len() {
269            return Err(io::Error::new(
270                io::ErrorKind::UnexpectedEof,
271                "cold index page ended early",
272            ));
273        }
274        let slice = &self.bytes[self.offset..end];
275        self.offset = end;
276        Ok(slice)
277    }
278
279    fn read_u8(&mut self) -> io::Result<u8> {
280        Ok(self.read_exact(1)?[0])
281    }
282
283    fn read_u16(&mut self) -> io::Result<u16> {
284        let mut bytes = [0; 2];
285        bytes.copy_from_slice(self.read_exact(2)?);
286        Ok(u16::from_le_bytes(bytes))
287    }
288
289    fn read_u32(&mut self) -> io::Result<u32> {
290        let mut bytes = [0; 4];
291        bytes.copy_from_slice(self.read_exact(4)?);
292        Ok(u32::from_le_bytes(bytes))
293    }
294
295    fn read_u64(&mut self) -> io::Result<u64> {
296        let mut bytes = [0; 8];
297        bytes.copy_from_slice(self.read_exact(8)?);
298        Ok(u64::from_le_bytes(bytes))
299    }
300
301    fn read_string(&mut self) -> io::Result<String> {
302        let len = usize::try_from(self.read_u32()?).expect("u32 fits usize");
303        let bytes = self.read_exact(len)?;
304        String::from_utf8(bytes.to_vec()).map_err(|err| {
305            io::Error::new(
306                io::ErrorKind::InvalidData,
307                format!("cold index page contains invalid UTF-8: {err}"),
308            )
309        })
310    }
311}
312
313pub trait ColdIndexPageStore: Send + Sync {
314    fn put_page<'a>(
315        &'a self,
316        key: &'a ColdIndexPageKey,
317        page: &'a ColdIndexPage,
318    ) -> ColdIndexPageStoreFuture<'a, ()>;
319
320    fn get_page<'a>(
321        &'a self,
322        key: &'a ColdIndexPageKey,
323    ) -> ColdIndexPageStoreFuture<'a, Option<ColdIndexPage>>;
324}
325
326pub async fn write_cold_chunk_index_pages<S: ColdIndexPageStore + ?Sized>(
327    store: &S,
328    stream_id: &BucketStreamId,
329    chunk: &ColdChunkRef,
330) -> io::Result<()> {
331    write_cold_chunk_index_pages_with_rollback(store, stream_id, chunk)
332        .await
333        .map(|_| ())
334}
335
336pub async fn write_cold_chunk_index_pages_with_rollback<S: ColdIndexPageStore + ?Sized>(
337    store: &S,
338    stream_id: &BucketStreamId,
339    chunk: &ColdChunkRef,
340) -> io::Result<Vec<ColdIndexPageRollback>> {
341    if chunk.end_offset <= chunk.start_offset {
342        return Ok(Vec::new());
343    }
344    let first_page_id = chunk.start_offset / ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES;
345    let last_page_id = (chunk.end_offset - 1) / ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES;
346    let mut rollback = Vec::new();
347    for page_id in first_page_id..=last_page_id {
348        let key = ColdIndexPageKey {
349            stream_id: stream_id.clone(),
350            generation: 0,
351            page_id,
352        };
353        let page_start = page_id.saturating_mul(ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES);
354        let page_end = page_start.saturating_add(ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES);
355        let previous = store.get_page(&key).await?;
356        let mut page = previous.clone().unwrap_or_else(|| ColdIndexPage {
357            start_offset: page_start,
358            end_offset: page_end,
359            cold_chunks: Vec::new(),
360            external_segments: Vec::new(),
361        });
362        rollback.push(ColdIndexPageRollback {
363            key: key.clone(),
364            previous,
365            written_chunk: chunk.clone(),
366        });
367        page.cold_chunks.retain(|existing| {
368            existing.start_offset != chunk.start_offset || existing.end_offset != chunk.end_offset
369        });
370        page.cold_chunks.push(chunk.clone());
371        page.cold_chunks.sort_by_key(|chunk| chunk.start_offset);
372        store.put_page(&key, &page).await?;
373    }
374    Ok(rollback)
375}
376
377pub async fn rollback_cold_index_pages<S: ColdIndexPageStore + ?Sized>(
378    store: &S,
379    rollback: Vec<ColdIndexPageRollback>,
380) -> io::Result<()> {
381    for entry in rollback.into_iter().rev() {
382        let Some(current) = store.get_page(&entry.key).await? else {
383            continue;
384        };
385        let current_still_has_written_chunk = current.cold_chunks.iter().any(|chunk| {
386            chunk.start_offset == entry.written_chunk.start_offset
387                && chunk.end_offset == entry.written_chunk.end_offset
388                && chunk.s3_path == entry.written_chunk.s3_path
389        });
390        if !current_still_has_written_chunk {
391            continue;
392        }
393        let page = entry.previous.unwrap_or_else(|| {
394            let page_start = entry
395                .key
396                .page_id
397                .saturating_mul(ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES);
398            let page_end = page_start.saturating_add(ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES);
399            ColdIndexPage {
400                start_offset: page_start,
401                end_offset: page_end,
402                cold_chunks: Vec::new(),
403                external_segments: Vec::new(),
404            }
405        });
406        store.put_page(&entry.key, &page).await?;
407    }
408    Ok(())
409}
410
411pub async fn write_external_segment_index_pages<S: ColdIndexPageStore + ?Sized>(
412    store: &S,
413    stream_id: &BucketStreamId,
414    start_offset: u64,
415    payload: &ExternalPayloadRef,
416) -> io::Result<()> {
417    let object = ObjectPayloadRef {
418        start_offset,
419        end_offset: start_offset.saturating_add(payload.payload_len),
420        s3_path: payload.s3_path.clone(),
421        object_size: payload.object_size,
422    };
423    write_object_index_pages(store, stream_id, object).await
424}
425
426async fn write_object_index_pages<S: ColdIndexPageStore + ?Sized>(
427    store: &S,
428    stream_id: &BucketStreamId,
429    object: ObjectPayloadRef,
430) -> io::Result<()> {
431    if object.end_offset <= object.start_offset {
432        return Ok(());
433    }
434    let first_page_id = object.start_offset / ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES;
435    let last_page_id = (object.end_offset - 1) / ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES;
436    for page_id in first_page_id..=last_page_id {
437        let key = ColdIndexPageKey {
438            stream_id: stream_id.clone(),
439            generation: 0,
440            page_id,
441        };
442        let page_start = page_id.saturating_mul(ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES);
443        let page_end = page_start.saturating_add(ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES);
444        let mut page = store
445            .get_page(&key)
446            .await?
447            .unwrap_or_else(|| ColdIndexPage {
448                start_offset: page_start,
449                end_offset: page_end,
450                cold_chunks: Vec::new(),
451                external_segments: Vec::new(),
452            });
453        page.external_segments.retain(|existing| {
454            existing.start_offset != object.start_offset || existing.end_offset != object.end_offset
455        });
456        page.external_segments.push(object.clone());
457        page.external_segments
458            .sort_by_key(|object| object.start_offset);
459        store.put_page(&key, &page).await?;
460    }
461    Ok(())
462}
463
464#[derive(Debug, Default)]
465pub struct InMemoryColdIndexPageStore {
466    pages: Mutex<HashMap<ColdIndexPageKey, Vec<u8>>>,
467}
468
469#[derive(Debug, Clone)]
470pub struct ColdStoreColdIndexPageStore {
471    cold_store: ColdStoreHandle,
472}
473
474impl ColdStoreColdIndexPageStore {
475    pub fn new(cold_store: ColdStoreHandle) -> Self {
476        Self { cold_store }
477    }
478}
479
480impl ColdIndexPageStore for ColdStoreColdIndexPageStore {
481    fn put_page<'a>(
482        &'a self,
483        key: &'a ColdIndexPageKey,
484        page: &'a ColdIndexPage,
485    ) -> ColdIndexPageStoreFuture<'a, ()> {
486        Box::pin(async move {
487            let bytes = encode_page(key, page);
488            self.cold_store
489                .write_cold_index_page(&key.path(), &bytes)
490                .await?;
491            Ok(())
492        })
493    }
494
495    fn get_page<'a>(
496        &'a self,
497        key: &'a ColdIndexPageKey,
498    ) -> ColdIndexPageStoreFuture<'a, Option<ColdIndexPage>> {
499        Box::pin(async move {
500            self.cold_store
501                .read_cold_index_page(&key.path())
502                .await?
503                .map(|bytes| decode_page(key, &bytes))
504                .transpose()
505        })
506    }
507}
508
509impl InMemoryColdIndexPageStore {
510    pub fn new() -> Self {
511        Self::default()
512    }
513}
514
515impl ColdIndexPageStore for InMemoryColdIndexPageStore {
516    fn put_page<'a>(
517        &'a self,
518        key: &'a ColdIndexPageKey,
519        page: &'a ColdIndexPage,
520    ) -> ColdIndexPageStoreFuture<'a, ()> {
521        Box::pin(async move {
522            let bytes = encode_page(key, page);
523            self.pages
524                .lock()
525                .expect("cold index page store mutex poisoned")
526                .insert(key.clone(), bytes);
527            Ok(())
528        })
529    }
530
531    fn get_page<'a>(
532        &'a self,
533        key: &'a ColdIndexPageKey,
534    ) -> ColdIndexPageStoreFuture<'a, Option<ColdIndexPage>> {
535        Box::pin(async move {
536            self.pages
537                .lock()
538                .expect("cold index page store mutex poisoned")
539                .get(key)
540                .map(|bytes| decode_page(key, bytes))
541                .transpose()
542        })
543    }
544}
545
546#[derive(Debug)]
547pub struct ColdIndexPageCache<S: ColdIndexPageStore + ?Sized> {
548    store: Arc<S>,
549    capacity_pages: usize,
550    inner: Mutex<ColdIndexPageCacheInner>,
551}
552
553#[derive(Debug, Default)]
554struct ColdIndexPageCacheInner {
555    next_generation: u64,
556    pages: HashMap<ColdIndexPageKey, ColdIndexPageCacheEntry>,
557    lru: VecDeque<(ColdIndexPageKey, u64)>,
558}
559
560#[derive(Debug)]
561struct ColdIndexPageCacheEntry {
562    page: Arc<ColdIndexPage>,
563    generation: u64,
564}
565
566impl<S: ColdIndexPageStore + ?Sized> ColdIndexPageCache<S> {
567    pub fn new(store: Arc<S>, capacity_pages: usize) -> Self {
568        Self {
569            store,
570            capacity_pages,
571            inner: Mutex::new(ColdIndexPageCacheInner::default()),
572        }
573    }
574
575    pub async fn put_page(&self, key: &ColdIndexPageKey, page: &ColdIndexPage) -> io::Result<()> {
576        self.store.put_page(key, page).await?;
577        self.insert(key.clone(), Arc::new(page.clone()));
578        Ok(())
579    }
580
581    pub async fn get_page(&self, key: &ColdIndexPageKey) -> io::Result<Option<Arc<ColdIndexPage>>> {
582        if let Some(page) = self.get_cached(key) {
583            return Ok(Some(page));
584        }
585        self.reload_page(key).await
586    }
587
588    async fn reload_page(&self, key: &ColdIndexPageKey) -> io::Result<Option<Arc<ColdIndexPage>>> {
589        let Some(page) = self.store.get_page(key).await? else {
590            return Ok(None);
591        };
592        let page = Arc::new(page);
593        self.insert(key.clone(), page.clone());
594        Ok(Some(page))
595    }
596
597    pub async fn object_segments_for_read(
598        &self,
599        stream_id: &BucketStreamId,
600        segment: &StreamReadColdIndexSegment,
601    ) -> io::Result<Vec<ObjectPayloadRef>> {
602        let key = ColdIndexPageKey {
603            stream_id: stream_id.clone(),
604            generation: segment.generation,
605            page_id: segment.page_id,
606        };
607        let Some(page) = self.get_page(&key).await? else {
608            return Err(io::Error::new(
609                io::ErrorKind::NotFound,
610                format!("cold index page '{}' does not exist", key.path()),
611            ));
612        };
613        let read_end = segment
614            .read_start_offset
615            .checked_add(u64::try_from(segment.len).expect("cold index read len fits u64"))
616            .ok_or_else(|| {
617                io::Error::new(
618                    io::ErrorKind::InvalidInput,
619                    "cold index read range overflows",
620                )
621            })?;
622        let mut objects = objects_for_read(&page, segment.read_start_offset, read_end);
623        if !objects_cover_range(&objects, segment.read_start_offset, read_end)
624            && let Some(reloaded) = self.reload_page(&key).await?
625        {
626            objects = objects_for_read(&reloaded, segment.read_start_offset, read_end);
627        }
628        if !objects_cover_range(&objects, segment.read_start_offset, read_end) {
629            return Err(io::Error::new(
630                io::ErrorKind::InvalidData,
631                "cold index page does not cover requested read range",
632            ));
633        }
634        Ok(objects)
635    }
636
637    pub fn cached_page_count(&self) -> usize {
638        self.inner
639            .lock()
640            .expect("cold index page cache mutex poisoned")
641            .pages
642            .len()
643    }
644
645    fn get_cached(&self, key: &ColdIndexPageKey) -> Option<Arc<ColdIndexPage>> {
646        let mut inner = self
647            .inner
648            .lock()
649            .expect("cold index page cache mutex poisoned");
650        let page = inner.pages.get(key)?.page.clone();
651        Self::touch(&mut inner, key.clone());
652        Some(page)
653    }
654
655    fn insert(&self, key: ColdIndexPageKey, page: Arc<ColdIndexPage>) {
656        let mut inner = self
657            .inner
658            .lock()
659            .expect("cold index page cache mutex poisoned");
660        let generation = Self::touch(&mut inner, key.clone());
661        inner
662            .pages
663            .insert(key, ColdIndexPageCacheEntry { page, generation });
664        Self::evict_over_capacity(&mut inner, self.capacity_pages);
665    }
666
667    fn touch(inner: &mut ColdIndexPageCacheInner, key: ColdIndexPageKey) -> u64 {
668        let generation = inner.next_generation;
669        inner.next_generation = inner.next_generation.saturating_add(1);
670        if let Some(entry) = inner.pages.get_mut(&key) {
671            entry.generation = generation;
672        }
673        inner.lru.push_back((key, generation));
674        generation
675    }
676
677    fn evict_over_capacity(inner: &mut ColdIndexPageCacheInner, capacity_pages: usize) {
678        if capacity_pages == 0 {
679            inner.pages.clear();
680            inner.lru.clear();
681            return;
682        }
683        while inner.pages.len() > capacity_pages {
684            let Some((key, generation)) = inner.lru.pop_front() else {
685                break;
686            };
687            let stale = inner
688                .pages
689                .get(&key)
690                .is_none_or(|entry| entry.generation != generation);
691            if stale {
692                continue;
693            }
694            inner.pages.remove(&key);
695        }
696    }
697}
698
699fn objects_for_read(page: &ColdIndexPage, read_start: u64, read_end: u64) -> Vec<ObjectPayloadRef> {
700    let mut objects = Vec::new();
701    for chunk in &page.cold_chunks {
702        if let Some(object) = intersect_object(&ObjectPayloadRef::from(chunk), read_start, read_end)
703        {
704            objects.push(object);
705        }
706    }
707    for object in &page.external_segments {
708        if let Some(object) = intersect_object(object, read_start, read_end) {
709            objects.push(object);
710        }
711    }
712    objects.sort_by_key(|object| object.start_offset);
713    objects
714}
715
716fn intersect_object(
717    object: &ObjectPayloadRef,
718    read_start: u64,
719    read_end: u64,
720) -> Option<ObjectPayloadRef> {
721    let start = object.start_offset.max(read_start);
722    let end = object.end_offset.min(read_end);
723    (start < end).then(|| object.clone())
724}
725
726fn objects_cover_range(objects: &[ObjectPayloadRef], start: u64, end: u64) -> bool {
727    let mut expected = start;
728    for object in objects {
729        if object.end_offset <= expected {
730            continue;
731        }
732        if object.start_offset > expected {
733            return false;
734        }
735        expected = object.end_offset;
736        if expected >= end {
737            return true;
738        }
739    }
740    expected == end
741}
742
743#[cfg(test)]
744mod tests {
745    use super::*;
746
747    fn key(page_id: u64) -> ColdIndexPageKey {
748        ColdIndexPageKey {
749            stream_id: BucketStreamId::new("benchcmp", "cold-index"),
750            generation: 7,
751            page_id,
752        }
753    }
754
755    fn page(start_offset: u64, end_offset: u64) -> ColdIndexPage {
756        ColdIndexPage {
757            start_offset,
758            end_offset,
759            cold_chunks: vec![ColdChunkRef {
760                start_offset,
761                end_offset,
762                s3_path: format!("benchcmp/cold-index/chunks/{start_offset:020}.bin"),
763                object_size: end_offset - start_offset,
764            }],
765            external_segments: Vec::new(),
766        }
767    }
768
769    #[tokio::test]
770    async fn memory_store_round_trips_pages() {
771        let store = InMemoryColdIndexPageStore::new();
772        let key = key(1);
773        let page = page(0, 128);
774
775        assert_eq!(
776            key.path(),
777            "benchcmp/cold-index/cold-index/00000000000000000007/00000000000000000001.idx"
778        );
779        assert_eq!(store.get_page(&key).await.expect("get missing"), None);
780        store.put_page(&key, &page).await.expect("put page");
781        assert_eq!(
782            store.get_page(&key).await.expect("get page"),
783            Some(page.clone())
784        );
785        assert!(page.covers(127));
786        assert!(!page.covers(128));
787    }
788
789    #[tokio::test]
790    async fn rollback_skips_page_updated_by_newer_writer() {
791        let store = InMemoryColdIndexPageStore::new();
792        let stream_id = BucketStreamId::new("benchcmp", "cold-index");
793        let first = ColdChunkRef {
794            start_offset: 0,
795            end_offset: 128,
796            s3_path: "benchcmp/cold-index/chunks/first.bin".to_owned(),
797            object_size: 128,
798        };
799        let stale = ColdChunkRef {
800            start_offset: 0,
801            end_offset: 128,
802            s3_path: "benchcmp/cold-index/chunks/stale.bin".to_owned(),
803            object_size: 128,
804        };
805        let newer = ColdChunkRef {
806            start_offset: 0,
807            end_offset: 128,
808            s3_path: "benchcmp/cold-index/chunks/newer.bin".to_owned(),
809            object_size: 128,
810        };
811        write_cold_chunk_index_pages(&store, &stream_id, &first)
812            .await
813            .expect("write first chunk");
814        let rollback = write_cold_chunk_index_pages_with_rollback(&store, &stream_id, &stale)
815            .await
816            .expect("write stale chunk");
817        write_cold_chunk_index_pages(&store, &stream_id, &newer)
818            .await
819            .expect("write newer chunk");
820
821        rollback_cold_index_pages(&store, rollback)
822            .await
823            .expect("rollback stale chunk");
824
825        let page = store
826            .get_page(&ColdIndexPageKey {
827                stream_id,
828                generation: 0,
829                page_id: 0,
830            })
831            .await
832            .expect("get page")
833            .expect("page exists");
834        assert_eq!(page.cold_chunks, vec![newer]);
835    }
836
837    #[tokio::test]
838    async fn read_reload_repairs_stale_cached_page() {
839        let store = Arc::new(InMemoryColdIndexPageStore::new());
840        let stream_id = BucketStreamId::new("benchcmp", "cold-index");
841        let cache = ColdIndexPageCache::new(store.clone(), 8);
842        let first = ColdChunkRef {
843            start_offset: 0,
844            end_offset: 128,
845            s3_path: "benchcmp/cold-index/chunks/first.bin".to_owned(),
846            object_size: 128,
847        };
848        write_cold_chunk_index_pages(store.as_ref(), &stream_id, &first)
849            .await
850            .expect("write first chunk");
851        assert_eq!(
852            cache
853                .object_segments_for_read(&stream_id, &StreamReadColdIndexSegment {
854                    generation: 0,
855                    page_id: 0,
856                    read_start_offset: 0,
857                    len: 1,
858                },)
859                .await
860                .expect("read first byte")
861                .len(),
862            1
863        );
864
865        let second = ColdChunkRef {
866            start_offset: 128,
867            end_offset: 256,
868            s3_path: "benchcmp/cold-index/chunks/second.bin".to_owned(),
869            object_size: 128,
870        };
871        write_cold_chunk_index_pages(store.as_ref(), &stream_id, &second)
872            .await
873            .expect("write second chunk behind cache");
874
875        let objects = cache
876            .object_segments_for_read(&stream_id, &StreamReadColdIndexSegment {
877                generation: 0,
878                page_id: 0,
879                read_start_offset: 128,
880                len: 1,
881            })
882            .await
883            .expect("reload stale page");
884        assert_eq!(objects[0].s3_path, second.s3_path);
885    }
886
887    #[test]
888    fn binary_page_format_round_trips_and_validates() {
889        let key = key(42);
890        let mut page = page(128, 256);
891        page.external_segments.push(ObjectPayloadRef {
892            start_offset: 256,
893            end_offset: 300,
894            s3_path: "benchcmp/cold-index/external/00000000000000000256.bin".to_owned(),
895            object_size: 44,
896        });
897        let bytes = encode_page(&key, &page);
898        assert!(bytes.starts_with(COLD_INDEX_PAGE_MAGIC));
899
900        assert_eq!(decode_page(&key, &bytes).expect("decode page"), page);
901
902        let mut corrupted = bytes.clone();
903        let last = corrupted.last_mut().expect("checksum byte");
904        *last ^= 0xff;
905        let err = decode_page(&key, &corrupted).expect_err("corrupt checksum");
906        assert_eq!(err.kind(), io::ErrorKind::InvalidData);
907
908        let wrong_key = ColdIndexPageKey {
909            stream_id: key.stream_id.clone(),
910            generation: key.generation + 1,
911            page_id: key.page_id,
912        };
913        let err = decode_page(&wrong_key, &bytes).expect_err("key mismatch");
914        assert_eq!(err.kind(), io::ErrorKind::InvalidData);
915    }
916
917    #[tokio::test]
918    async fn page_cache_loads_on_miss_and_evicts_lru() {
919        let store = Arc::new(InMemoryColdIndexPageStore::new());
920        for page_id in 0..3 {
921            store
922                .put_page(&key(page_id), &page(page_id * 100, page_id * 100 + 100))
923                .await
924                .expect("put page");
925        }
926        let cache = ColdIndexPageCache::new(store, 2);
927
928        assert_eq!(
929            cache
930                .get_page(&key(0))
931                .await
932                .expect("load page")
933                .expect("page")
934                .start_offset,
935            0
936        );
937        assert_eq!(
938            cache
939                .get_page(&key(1))
940                .await
941                .expect("load page")
942                .expect("page")
943                .start_offset,
944            100
945        );
946        assert_eq!(cache.cached_page_count(), 2);
947
948        // Touch page 0 so page 1 becomes the eviction candidate.
949        assert!(
950            cache
951                .get_page(&key(0))
952                .await
953                .expect("cached page")
954                .is_some()
955        );
956        assert_eq!(
957            cache
958                .get_page(&key(2))
959                .await
960                .expect("load page")
961                .expect("page")
962                .start_offset,
963            200
964        );
965        assert_eq!(cache.cached_page_count(), 2);
966    }
967
968    #[tokio::test]
969    async fn zero_capacity_cache_does_not_retain_pages() {
970        let store = Arc::new(InMemoryColdIndexPageStore::new());
971        store
972            .put_page(&key(0), &page(0, 64))
973            .await
974            .expect("put page");
975        let cache = ColdIndexPageCache::new(store, 0);
976
977        assert!(cache.get_page(&key(0)).await.expect("load page").is_some());
978        assert_eq!(cache.cached_page_count(), 0);
979    }
980}
ursula_runtime/cold_index.rs

ursula_runtime/
cold_index.rs
