ursula_runtime/

cold_index.rs

use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::VecDeque;
use std::future::Future;
use std::io;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::Mutex;

use ursula_shard::BucketStreamId;
use ursula_stream::ColdChunkRef;
use ursula_stream::ExternalPayloadRef;
use ursula_stream::ObjectPayloadRef;
use ursula_stream::StreamReadColdIndexSegment;

use crate::cold_store::ColdStoreHandle;

pub type ColdIndexPageStoreFuture<'a, T> = Pin<Box<dyn Future<Output = io::Result<T>> + Send + 'a>>;

const COLD_INDEX_PAGE_MAGIC: &[u8; 8] = b"UCIDX001";
const COLD_INDEX_PAGE_VERSION: u16 = 2;
const COLD_INDEX_ENTRY_COLD_CHUNK: u8 = 1;
const COLD_INDEX_ENTRY_EXTERNAL_SEGMENT: u8 = 2;
const FNV64_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
const FNV64_PRIME: u64 = 0x0000_0100_0000_01b3;

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ColdIndexPageKey {
    pub stream_id: BucketStreamId,
    pub generation: u64,
    pub page_id: u64,
}

impl ColdIndexPageKey {
    pub fn path(&self) -> String {
        format!(
            "{}/cold-index/{:020}/{:020}.idx",
            self.stream_id, self.generation, self.page_id
        )
    }
}

pub fn cold_index_prefix(stream_id: &BucketStreamId) -> String {
    format!("{stream_id}/cold-index/")
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ColdIndexPage {
    pub start_offset: u64,
    pub end_offset: u64,
    pub cold_chunks: Vec<ColdChunkRef>,
    pub external_segments: Vec<ObjectPayloadRef>,
}

impl ColdIndexPage {
    pub fn covers(&self, offset: u64) -> bool {
        self.start_offset <= offset && offset < self.end_offset
    }
}

#[derive(Debug, Clone)]
pub struct ColdIndexPageRollback {
    key: ColdIndexPageKey,
    previous: Option<ColdIndexPage>,
    written_chunk: ColdChunkRef,
}

fn encode_page(key: &ColdIndexPageKey, page: &ColdIndexPage) -> Vec<u8> {
    let mut body = Vec::new();
    put_string(&mut body, &key.stream_id.bucket_id);
    match &key.stream_id.affinity_key {
        Some(affinity_key) => {
            put_u8(&mut body, 1);
            put_string(&mut body, affinity_key);
        }
        None => put_u8(&mut body, 0),
    }
    put_string(&mut body, &key.stream_id.stream_id);
    put_u64(&mut body, key.generation);
    put_u64(&mut body, key.page_id);
    put_u64(&mut body, page.start_offset);
    put_u64(&mut body, page.end_offset);
    put_u32(
        &mut body,
        u32::try_from(page.cold_chunks.len()).expect("cold index cold chunk count fits u32"),
    );
    for chunk in &page.cold_chunks {
        put_u8(&mut body, COLD_INDEX_ENTRY_COLD_CHUNK);
        put_u64(&mut body, chunk.start_offset);
        put_u64(&mut body, chunk.end_offset);
        put_u64(&mut body, chunk.object_size);
        put_string(&mut body, &chunk.s3_path);
    }
    put_u32(
        &mut body,
        u32::try_from(page.external_segments.len())
            .expect("cold index external segment count fits u32"),
    );
    for object in &page.external_segments {
        put_u8(&mut body, COLD_INDEX_ENTRY_EXTERNAL_SEGMENT);
        put_u64(&mut body, object.start_offset);
        put_u64(&mut body, object.end_offset);
        put_u64(&mut body, object.object_size);
        put_string(&mut body, &object.s3_path);
    }

    let mut bytes = Vec::with_capacity(COLD_INDEX_PAGE_MAGIC.len() + 2 + 4 + body.len() + 8);
    bytes.extend_from_slice(COLD_INDEX_PAGE_MAGIC);
    put_u16(&mut bytes, COLD_INDEX_PAGE_VERSION);
    put_u32(
        &mut bytes,
        u32::try_from(body.len()).expect("cold index page body len fits u32"),
    );
    bytes.extend_from_slice(&body);
    put_u64(&mut bytes, checksum64(&body));
    bytes
}

fn decode_page(key: &ColdIndexPageKey, bytes: &[u8]) -> io::Result<ColdIndexPage> {
    let mut cursor = Cursor::new(bytes);
    let magic = cursor.read_exact(COLD_INDEX_PAGE_MAGIC.len())?;
    if magic != COLD_INDEX_PAGE_MAGIC {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "cold index page has invalid magic",
        ));
    }
    let version = cursor.read_u16()?;
    if !matches!(version, 1 | COLD_INDEX_PAGE_VERSION) {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("unsupported cold index page version {version}"),
        ));
    }
    let body_len = usize::try_from(cursor.read_u32()?).expect("u32 fits usize");
    let body = cursor.read_exact(body_len)?;
    let expected_checksum = cursor.read_u64()?;
    if cursor.remaining() != 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "cold index page has trailing bytes",
        ));
    }
    let actual_checksum = checksum64(body);
    if actual_checksum != expected_checksum {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "cold index page checksum mismatch",
        ));
    }

    let mut body = Cursor::new(body);
    let bucket_id = body.read_string()?;
    let affinity_key = if version >= 2 {
        match body.read_u8()? {
            0 => None,
            1 => Some(body.read_string()?),
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "cold index page has invalid affinity marker",
                ));
            }
        }
    } else {
        None
    };
    let stream_id = body.read_string()?;
    let generation = body.read_u64()?;
    let page_id = body.read_u64()?;
    if bucket_id != key.stream_id.bucket_id
        || affinity_key != key.stream_id.affinity_key
        || stream_id != key.stream_id.stream_id
        || generation != key.generation
        || page_id != key.page_id
    {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "cold index page key metadata mismatch",
        ));
    }
    let start_offset = body.read_u64()?;
    let end_offset = body.read_u64()?;
    let cold_chunk_count = body.read_u32()?;
    let mut cold_chunks =
        Vec::with_capacity(usize::try_from(cold_chunk_count).expect("u32 fits usize"));
    for _ in 0..cold_chunk_count {
        let tag = body.read_u8()?;
        if tag != COLD_INDEX_ENTRY_COLD_CHUNK {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "cold index page expected cold chunk entry",
            ));
        }
        cold_chunks.push(ColdChunkRef {
            start_offset: body.read_u64()?,
            end_offset: body.read_u64()?,
            object_size: body.read_u64()?,
            s3_path: body.read_string()?,
            object_offset: 0,
            shared_object: false,
            payload_digest: String::new(),
        });
    }
    let external_segment_count = body.read_u32()?;
    let mut external_segments =
        Vec::with_capacity(usize::try_from(external_segment_count).expect("u32 fits usize"));
    for _ in 0..external_segment_count {
        let tag = body.read_u8()?;
        if tag != COLD_INDEX_ENTRY_EXTERNAL_SEGMENT {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "cold index page expected external segment entry",
            ));
        }
        external_segments.push(ObjectPayloadRef {
            start_offset: body.read_u64()?,
            end_offset: body.read_u64()?,
            object_size: body.read_u64()?,
            s3_path: body.read_string()?,
            object_offset: 0,
        });
    }
    if body.remaining() != 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "cold index page body has trailing bytes",
        ));
    }
    Ok(ColdIndexPage {
        start_offset,
        end_offset,
        cold_chunks,
        external_segments,
    })
}

fn put_u8(out: &mut Vec<u8>, value: u8) {
    out.push(value);
}

fn put_u16(out: &mut Vec<u8>, value: u16) {
    out.extend_from_slice(&value.to_le_bytes());
}

fn put_u32(out: &mut Vec<u8>, value: u32) {
    out.extend_from_slice(&value.to_le_bytes());
}

fn put_u64(out: &mut Vec<u8>, value: u64) {
    out.extend_from_slice(&value.to_le_bytes());
}

fn put_string(out: &mut Vec<u8>, value: &str) {
    put_u32(
        out,
        u32::try_from(value.len()).expect("cold index string len fits u32"),
    );
    out.extend_from_slice(value.as_bytes());
}

fn checksum64(bytes: &[u8]) -> u64 {
    let mut hash = FNV64_OFFSET_BASIS;
    for byte in bytes {
        hash ^= u64::from(*byte);
        hash = hash.wrapping_mul(FNV64_PRIME);
    }
    hash
}

struct Cursor<'a> {
    bytes: &'a [u8],
    offset: usize,
}

impl<'a> Cursor<'a> {
    fn new(bytes: &'a [u8]) -> Self {
        Self { bytes, offset: 0 }
    }

    fn remaining(&self) -> usize {
        self.bytes.len().saturating_sub(self.offset)
    }

    fn read_exact(&mut self, len: usize) -> io::Result<&'a [u8]> {
        let end = self.offset.checked_add(len).ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidData,
                "cold index page offset overflow",
            )
        })?;
        if end > self.bytes.len() {
            return Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "cold index page ended early",
            ));
        }
        let slice = &self.bytes[self.offset..end];
        self.offset = end;
        Ok(slice)
    }

    fn read_u8(&mut self) -> io::Result<u8> {
        Ok(self.read_exact(1)?[0])
    }

    fn read_u16(&mut self) -> io::Result<u16> {
        let mut bytes = [0; 2];
        bytes.copy_from_slice(self.read_exact(2)?);
        Ok(u16::from_le_bytes(bytes))
    }

    fn read_u32(&mut self) -> io::Result<u32> {
        let mut bytes = [0; 4];
        bytes.copy_from_slice(self.read_exact(4)?);
        Ok(u32::from_le_bytes(bytes))
    }

    fn read_u64(&mut self) -> io::Result<u64> {
        let mut bytes = [0; 8];
        bytes.copy_from_slice(self.read_exact(8)?);
        Ok(u64::from_le_bytes(bytes))
    }

    fn read_string(&mut self) -> io::Result<String> {
        let len = usize::try_from(self.read_u32()?).expect("u32 fits usize");
        let bytes = self.read_exact(len)?;
        String::from_utf8(bytes.to_vec()).map_err(|err| {
            io::Error::new(
                io::ErrorKind::InvalidData,
                format!("cold index page contains invalid UTF-8: {err}"),
            )
        })
    }
}

pub trait ColdIndexPageStore: Send + Sync {
    fn put_page<'a>(
        &'a self,
        key: &'a ColdIndexPageKey,
        page: &'a ColdIndexPage,
    ) -> ColdIndexPageStoreFuture<'a, ()>;

    fn get_page<'a>(
        &'a self,
        key: &'a ColdIndexPageKey,
    ) -> ColdIndexPageStoreFuture<'a, Option<ColdIndexPage>>;
}

pub async fn write_cold_chunk_index_pages<S: ColdIndexPageStore + ?Sized>(
    store: &S,
    stream_id: &BucketStreamId,
    chunk: &ColdChunkRef,
) -> io::Result<()> {
    write_cold_chunk_index_pages_with_rollback(store, stream_id, chunk)
        .await
        .map(|_| ())
}

pub async fn write_cold_chunk_index_pages_with_rollback<S: ColdIndexPageStore + ?Sized>(
    store: &S,
    stream_id: &BucketStreamId,
    chunk: &ColdChunkRef,
) -> io::Result<Vec<ColdIndexPageRollback>> {
    if chunk.end_offset <= chunk.start_offset {
        return Ok(Vec::new());
    }
    let first_page_id = chunk.start_offset / ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES;
    let last_page_id = (chunk.end_offset - 1) / ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES;
    let mut rollback = Vec::new();
    for page_id in first_page_id..=last_page_id {
        let key = ColdIndexPageKey {
            stream_id: stream_id.clone(),
            generation: 0,
            page_id,
        };
        let page_start = page_id.saturating_mul(ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES);
        let page_end = page_start.saturating_add(ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES);
        let previous = store.get_page(&key).await?;
        let mut page = previous.clone().unwrap_or_else(|| ColdIndexPage {
            start_offset: page_start,
            end_offset: page_end,
            cold_chunks: Vec::new(),
            external_segments: Vec::new(),
        });
        rollback.push(ColdIndexPageRollback {
            key: key.clone(),
            previous,
            written_chunk: chunk.clone(),
        });
        page.cold_chunks.retain(|existing| {
            existing.start_offset != chunk.start_offset || existing.end_offset != chunk.end_offset
        });
        page.cold_chunks.push(chunk.clone());
        page.cold_chunks.sort_by_key(|chunk| chunk.start_offset);
        store.put_page(&key, &page).await?;
    }
    Ok(rollback)
}

pub async fn rollback_cold_index_pages<S: ColdIndexPageStore + ?Sized>(
    store: &S,
    rollback: Vec<ColdIndexPageRollback>,
) -> io::Result<()> {
    for entry in rollback.into_iter().rev() {
        let Some(current) = store.get_page(&entry.key).await? else {
            continue;
        };
        let current_still_has_written_chunk = current.cold_chunks.iter().any(|chunk| {
            chunk.start_offset == entry.written_chunk.start_offset
                && chunk.end_offset == entry.written_chunk.end_offset
                && chunk.s3_path == entry.written_chunk.s3_path
        });
        if !current_still_has_written_chunk {
            continue;
        }
        let page = entry.previous.unwrap_or_else(|| {
            let page_start = entry
                .key
                .page_id
                .saturating_mul(ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES);
            let page_end = page_start.saturating_add(ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES);
            ColdIndexPage {
                start_offset: page_start,
                end_offset: page_end,
                cold_chunks: Vec::new(),
                external_segments: Vec::new(),
            }
        });
        store.put_page(&entry.key, &page).await?;
    }
    Ok(())
}

pub async fn write_external_segment_index_pages<S: ColdIndexPageStore + ?Sized>(
    store: &S,
    stream_id: &BucketStreamId,
    start_offset: u64,
    payload: &ExternalPayloadRef,
) -> io::Result<()> {
    let object = ObjectPayloadRef {
        start_offset,
        end_offset: start_offset.saturating_add(payload.payload_len),
        s3_path: payload.s3_path.clone(),
        object_size: payload.object_size,
        object_offset: 0,
    };
    write_object_index_pages(store, stream_id, object).await
}

async fn write_object_index_pages<S: ColdIndexPageStore + ?Sized>(
    store: &S,
    stream_id: &BucketStreamId,
    object: ObjectPayloadRef,
) -> io::Result<()> {
    if object.end_offset <= object.start_offset {
        return Ok(());
    }
    let first_page_id = object.start_offset / ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES;
    let last_page_id = (object.end_offset - 1) / ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES;
    for page_id in first_page_id..=last_page_id {
        let key = ColdIndexPageKey {
            stream_id: stream_id.clone(),
            generation: 0,
            page_id,
        };
        let page_start = page_id.saturating_mul(ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES);
        let page_end = page_start.saturating_add(ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES);
        let mut page = store
            .get_page(&key)
            .await?
            .unwrap_or_else(|| ColdIndexPage {
                start_offset: page_start,
                end_offset: page_end,
                cold_chunks: Vec::new(),
                external_segments: Vec::new(),
            });
        page.external_segments.retain(|existing| {
            existing.start_offset != object.start_offset || existing.end_offset != object.end_offset
        });
        page.external_segments.push(object.clone());
        page.external_segments
            .sort_by_key(|object| object.start_offset);
        store.put_page(&key, &page).await?;
    }
    Ok(())
}

#[derive(Debug, Default)]
pub struct InMemoryColdIndexPageStore {
    pages: Mutex<HashMap<ColdIndexPageKey, Vec<u8>>>,
}

#[derive(Debug, Clone)]
pub struct ColdStoreColdIndexPageStore {
    cold_store: ColdStoreHandle,
}

impl ColdStoreColdIndexPageStore {
    pub fn new(cold_store: ColdStoreHandle) -> Self {
        Self { cold_store }
    }
}

impl ColdIndexPageStore for ColdStoreColdIndexPageStore {
    fn put_page<'a>(
        &'a self,
        key: &'a ColdIndexPageKey,
        page: &'a ColdIndexPage,
    ) -> ColdIndexPageStoreFuture<'a, ()> {
        Box::pin(async move {
            let bytes = encode_page(key, page);
            self.cold_store
                .write_cold_index_page(&key.path(), &bytes)
                .await?;
            Ok(())
        })
    }

    fn get_page<'a>(
        &'a self,
        key: &'a ColdIndexPageKey,
    ) -> ColdIndexPageStoreFuture<'a, Option<ColdIndexPage>> {
        Box::pin(async move {
            self.cold_store
                .read_cold_index_page(&key.path())
                .await?
                .map(|bytes| decode_page(key, &bytes))
                .transpose()
        })
    }
}

impl InMemoryColdIndexPageStore {
    pub fn new() -> Self {
        Self::default()
    }
}

impl ColdIndexPageStore for InMemoryColdIndexPageStore {
    fn put_page<'a>(
        &'a self,
        key: &'a ColdIndexPageKey,
        page: &'a ColdIndexPage,
    ) -> ColdIndexPageStoreFuture<'a, ()> {
        Box::pin(async move {
            let bytes = encode_page(key, page);
            self.pages
                .lock()
                .expect("cold index page store mutex poisoned")
                .insert(key.clone(), bytes);
            Ok(())
        })
    }

    fn get_page<'a>(
        &'a self,
        key: &'a ColdIndexPageKey,
    ) -> ColdIndexPageStoreFuture<'a, Option<ColdIndexPage>> {
        Box::pin(async move {
            self.pages
                .lock()
                .expect("cold index page store mutex poisoned")
                .get(key)
                .map(|bytes| decode_page(key, bytes))
                .transpose()
        })
    }
}

#[derive(Debug)]
pub struct ColdIndexPageCache<S: ColdIndexPageStore + ?Sized> {
    store: Arc<S>,
    capacity_pages: usize,
    inner: Mutex<ColdIndexPageCacheInner>,
}

#[derive(Debug, Default)]
struct ColdIndexPageCacheInner {
    next_generation: u64,
    pages: HashMap<ColdIndexPageKey, ColdIndexPageCacheEntry>,
    lru: VecDeque<(ColdIndexPageKey, u64)>,
}

#[derive(Debug)]
struct ColdIndexPageCacheEntry {
    page: Arc<ColdIndexPage>,
    generation: u64,
}

impl<S: ColdIndexPageStore + ?Sized> ColdIndexPageCache<S> {
    pub fn new(store: Arc<S>, capacity_pages: usize) -> Self {
        Self {
            store,
            capacity_pages,
            inner: Mutex::new(ColdIndexPageCacheInner::default()),
        }
    }

    pub async fn put_page(&self, key: &ColdIndexPageKey, page: &ColdIndexPage) -> io::Result<()> {
        self.store.put_page(key, page).await?;
        self.insert(key.clone(), Arc::new(page.clone()));
        Ok(())
    }

    pub async fn get_page(&self, key: &ColdIndexPageKey) -> io::Result<Option<Arc<ColdIndexPage>>> {
        if let Some(page) = self.get_cached(key) {
            return Ok(Some(page));
        }
        self.reload_page(key).await
    }

    /// Drops every cached generation/page for one stream. Compaction invokes
    /// this on every replica when the replicated replacement command applies.
    pub fn invalidate_stream(&self, stream_id: &BucketStreamId) {
        let mut inner = self.inner.lock().expect("cold index cache mutex poisoned");
        inner.pages.retain(|key, _| &key.stream_id != stream_id);
        inner.lru.retain(|(key, _)| &key.stream_id != stream_id);
    }

    async fn reload_page(&self, key: &ColdIndexPageKey) -> io::Result<Option<Arc<ColdIndexPage>>> {
        let Some(page) = self.store.get_page(key).await? else {
            return Ok(None);
        };
        let page = Arc::new(page);
        self.insert(key.clone(), page.clone());
        Ok(Some(page))
    }

    pub async fn object_segments_for_read(
        &self,
        stream_id: &BucketStreamId,
        segment: &StreamReadColdIndexSegment,
    ) -> io::Result<Vec<ObjectPayloadRef>> {
        let key = ColdIndexPageKey {
            stream_id: stream_id.clone(),
            generation: segment.generation,
            page_id: segment.page_id,
        };
        let Some(page) = self.get_page(&key).await? else {
            return Err(io::Error::new(
                io::ErrorKind::NotFound,
                format!("cold index page '{}' does not exist", key.path()),
            ));
        };
        let read_end = segment
            .read_start_offset
            .checked_add(u64::try_from(segment.len).expect("cold index read len fits u64"))
            .ok_or_else(|| {
                io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "cold index read range overflows",
                )
            })?;
        let mut objects = objects_for_read(&page, segment.read_start_offset, read_end);
        if !objects_cover_range(&objects, segment.read_start_offset, read_end)
            && let Some(reloaded) = self.reload_page(&key).await?
        {
            objects = objects_for_read(&reloaded, segment.read_start_offset, read_end);
        }
        if !objects_cover_range(&objects, segment.read_start_offset, read_end) {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "cold index page does not cover requested read range",
            ));
        }
        Ok(objects)
    }

    pub fn cached_page_count(&self) -> usize {
        self.inner
            .lock()
            .expect("cold index page cache mutex poisoned")
            .pages
            .len()
    }

    fn get_cached(&self, key: &ColdIndexPageKey) -> Option<Arc<ColdIndexPage>> {
        let mut inner = self
            .inner
            .lock()
            .expect("cold index page cache mutex poisoned");
        let page = inner.pages.get(key)?.page.clone();
        Self::touch(&mut inner, key.clone());
        Some(page)
    }

    fn insert(&self, key: ColdIndexPageKey, page: Arc<ColdIndexPage>) {
        let mut inner = self
            .inner
            .lock()
            .expect("cold index page cache mutex poisoned");
        let generation = Self::touch(&mut inner, key.clone());
        inner
            .pages
            .insert(key, ColdIndexPageCacheEntry { page, generation });
        Self::evict_over_capacity(&mut inner, self.capacity_pages);
    }

    fn touch(inner: &mut ColdIndexPageCacheInner, key: ColdIndexPageKey) -> u64 {
        let generation = inner.next_generation;
        inner.next_generation = inner.next_generation.saturating_add(1);
        if let Some(entry) = inner.pages.get_mut(&key) {
            entry.generation = generation;
        }
        inner.lru.push_back((key, generation));
        generation
    }

    fn evict_over_capacity(inner: &mut ColdIndexPageCacheInner, capacity_pages: usize) {
        if capacity_pages == 0 {
            inner.pages.clear();
            inner.lru.clear();
            return;
        }
        while inner.pages.len() > capacity_pages {
            let Some((key, generation)) = inner.lru.pop_front() else {
                break;
            };
            let stale = inner
                .pages
                .get(&key)
                .is_none_or(|entry| entry.generation != generation);
            if stale {
                continue;
            }
            inner.pages.remove(&key);
        }
    }
}

/// Loads and de-duplicates the chunk references present in a stream's index
/// pages. Chunks crossing a 64 MiB page boundary intentionally appear in more
/// than one page.
pub async fn load_cold_chunks_from_pages<S: ColdIndexPageStore + ?Sized>(
    store: &S,
    keys: &[ColdIndexPageKey],
) -> io::Result<Vec<ColdChunkRef>> {
    let mut chunks = Vec::new();
    let mut seen = HashSet::new();
    for key in keys {
        let Some(page) = store.get_page(key).await? else {
            continue;
        };
        for chunk in page.cold_chunks {
            let identity = (chunk.start_offset, chunk.end_offset, chunk.s3_path.clone());
            if seen.insert(identity) {
                chunks.push(chunk);
            }
        }
    }
    chunks.sort_by(|left, right| {
        left.start_offset
            .cmp(&right.start_offset)
            .then_with(|| left.end_offset.cmp(&right.end_offset))
            .then_with(|| left.s3_path.cmp(&right.s3_path))
    });
    Ok(chunks)
}

/// Selects the oldest contiguous run of undersized raw chunks whose combined
/// payload reaches the byte target without exceeding the configured maximum.
pub fn select_cold_chunk_compaction(
    chunks: &[ColdChunkRef],
    target_bytes: u64,
    max_bytes: u64,
) -> Option<Vec<ColdChunkRef>> {
    if target_bytes == 0 || max_bytes < target_bytes {
        return None;
    }
    let mut candidate = Vec::new();
    let mut bytes = 0_u64;
    let mut next_offset = None;
    for chunk in chunks {
        let logical_bytes = chunk.end_offset.checked_sub(chunk.start_offset)?;
        let usable = logical_bytes > 0
            && chunk.object_size == logical_bytes
            && chunk.object_size < target_bytes;
        let contiguous = next_offset.is_none_or(|offset| offset == chunk.start_offset);
        let next_bytes = bytes.checked_add(chunk.object_size);
        if !usable || !contiguous || next_bytes.is_none_or(|total| total > max_bytes) {
            candidate.clear();
            bytes = 0;
            next_offset = None;
            if !usable {
                continue;
            }
        }
        bytes = bytes.checked_add(chunk.object_size)?;
        next_offset = Some(chunk.end_offset);
        candidate.push(chunk.clone());
        if candidate.len() >= 2 && bytes >= target_bytes {
            return Some(candidate);
        }
    }
    None
}

/// Atomically at the page-object level replaces a contiguous set of chunk
/// references with one equivalent object. Every rewritten page always points
/// at readable old or new bytes, so a retry after a partial S3 failure remains
/// safe.
pub async fn replace_cold_chunk_index_pages<S: ColdIndexPageStore + ?Sized>(
    store: &S,
    stream_id: &BucketStreamId,
    old_chunks: &[ColdChunkRef],
    replacement: &ColdChunkRef,
) -> io::Result<bool> {
    replace_cold_chunk_index_pages_with_rollback(store, stream_id, old_chunks, replacement)
        .await
        .map(|rollback| rollback.is_some())
}

pub async fn replace_cold_chunk_index_pages_with_rollback<S: ColdIndexPageStore + ?Sized>(
    store: &S,
    stream_id: &BucketStreamId,
    old_chunks: &[ColdChunkRef],
    replacement: &ColdChunkRef,
) -> io::Result<Option<Vec<ColdIndexPageRollback>>> {
    if old_chunks.len() < 2 || replacement.end_offset <= replacement.start_offset {
        return Ok(None);
    }
    let first_page_id = replacement.start_offset / ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES;
    let last_page_id = (replacement.end_offset - 1) / ursula_stream::COLD_INDEX_PAGE_SPAN_BYTES;
    let old_identities = old_chunks
        .iter()
        .map(|chunk| (chunk.start_offset, chunk.end_offset, chunk.s3_path.as_str()))
        .collect::<HashSet<_>>();
    let mut pages = Vec::new();
    let mut found = HashSet::new();
    for page_id in first_page_id..=last_page_id {
        let key = ColdIndexPageKey {
            stream_id: stream_id.clone(),
            generation: 0,
            page_id,
        };
        let Some(mut page) = store.get_page(&key).await? else {
            return Ok(None);
        };
        let previous = page.clone();
        for chunk in &page.cold_chunks {
            let identity = (chunk.start_offset, chunk.end_offset, chunk.s3_path.as_str());
            if old_identities.contains(&identity) {
                found.insert((chunk.start_offset, chunk.end_offset, chunk.s3_path.clone()));
            }
        }
        page.cold_chunks.retain(|chunk| {
            !old_identities.contains(&(
                chunk.start_offset,
                chunk.end_offset,
                chunk.s3_path.as_str(),
            ))
        });
        page.cold_chunks.retain(|chunk| {
            chunk.start_offset != replacement.start_offset
                || chunk.end_offset != replacement.end_offset
        });
        page.cold_chunks.push(replacement.clone());
        page.cold_chunks.sort_by_key(|chunk| chunk.start_offset);
        pages.push((key, previous, page));
    }
    if found.len() != old_identities.len() {
        return Ok(None);
    }
    let mut rollback = Vec::with_capacity(pages.len());
    for (key, previous, page) in pages {
        if let Err(err) = store.put_page(&key, &page).await {
            rollback_cold_index_pages(store, rollback).await?;
            return Err(err);
        }
        rollback.push(ColdIndexPageRollback {
            key,
            previous: Some(previous),
            written_chunk: replacement.clone(),
        });
    }
    Ok(Some(rollback))
}

fn objects_for_read(page: &ColdIndexPage, read_start: u64, read_end: u64) -> Vec<ObjectPayloadRef> {
    let mut objects = Vec::new();
    for chunk in &page.cold_chunks {
        if let Some(object) = intersect_object(&ObjectPayloadRef::from(chunk), read_start, read_end)
        {
            objects.push(object);
        }
    }
    for object in &page.external_segments {
        if let Some(object) = intersect_object(object, read_start, read_end) {
            objects.push(object);
        }
    }
    objects.sort_by_key(|object| object.start_offset);
    objects
}

fn intersect_object(
    object: &ObjectPayloadRef,
    read_start: u64,
    read_end: u64,
) -> Option<ObjectPayloadRef> {
    let start = object.start_offset.max(read_start);
    let end = object.end_offset.min(read_end);
    (start < end).then(|| object.clone())
}

fn objects_cover_range(objects: &[ObjectPayloadRef], start: u64, end: u64) -> bool {
    let mut expected = start;
    for object in objects {
        if object.end_offset <= expected {
            continue;
        }
        if object.start_offset > expected {
            return false;
        }
        expected = object.end_offset;
        if expected >= end {
            return true;
        }
    }
    expected == end
}

#[cfg(test)]
mod tests {
    use super::*;

    fn key(page_id: u64) -> ColdIndexPageKey {
        ColdIndexPageKey {
            stream_id: BucketStreamId::new("benchcmp", "cold-index"),
            generation: 7,
            page_id,
        }
    }

    fn page(start_offset: u64, end_offset: u64) -> ColdIndexPage {
        ColdIndexPage {
            start_offset,
            end_offset,
            cold_chunks: vec![ColdChunkRef {
                start_offset,
                end_offset,
                s3_path: format!("benchcmp/cold-index/chunks/{start_offset:020}.bin"),
                object_size: end_offset - start_offset,
                ..Default::default()
            }],
            external_segments: Vec::new(),
        }
    }

    #[tokio::test]
    async fn memory_store_round_trips_pages() {
        let store = InMemoryColdIndexPageStore::new();
        let key = key(1);
        let page = page(0, 128);

        assert_eq!(
            key.path(),
            "benchcmp/cold-index/cold-index/00000000000000000007/00000000000000000001.idx"
        );
        assert_eq!(store.get_page(&key).await.expect("get missing"), None);
        store.put_page(&key, &page).await.expect("put page");
        assert_eq!(
            store.get_page(&key).await.expect("get page"),
            Some(page.clone())
        );
        assert!(page.covers(127));
        assert!(!page.covers(128));
    }

    #[tokio::test]
    async fn rollback_skips_page_updated_by_newer_writer() {
        let store = InMemoryColdIndexPageStore::new();
        let stream_id = BucketStreamId::new("benchcmp", "cold-index");
        let first = ColdChunkRef {
            start_offset: 0,
            end_offset: 128,
            s3_path: "benchcmp/cold-index/chunks/first.bin".to_owned(),
            object_size: 128,
            ..Default::default()
        };
        let stale = ColdChunkRef {
            start_offset: 0,
            end_offset: 128,
            s3_path: "benchcmp/cold-index/chunks/stale.bin".to_owned(),
            object_size: 128,
            ..Default::default()
        };
        let newer = ColdChunkRef {
            start_offset: 0,
            end_offset: 128,
            s3_path: "benchcmp/cold-index/chunks/newer.bin".to_owned(),
            object_size: 128,
            ..Default::default()
        };
        write_cold_chunk_index_pages(&store, &stream_id, &first)
            .await
            .expect("write first chunk");
        let rollback = write_cold_chunk_index_pages_with_rollback(&store, &stream_id, &stale)
            .await
            .expect("write stale chunk");
        write_cold_chunk_index_pages(&store, &stream_id, &newer)
            .await
            .expect("write newer chunk");

        rollback_cold_index_pages(&store, rollback)
            .await
            .expect("rollback stale chunk");

        let page = store
            .get_page(&ColdIndexPageKey {
                stream_id,
                generation: 0,
                page_id: 0,
            })
            .await
            .expect("get page")
            .expect("page exists");
        assert_eq!(page.cold_chunks, vec![newer]);
    }

    #[tokio::test]
    async fn compact_replacement_rollback_restores_input_chunks() {
        let store = InMemoryColdIndexPageStore::new();
        let stream_id = BucketStreamId::new("benchcmp", "cold-index");
        let first = ColdChunkRef {
            start_offset: 0,
            end_offset: 64,
            s3_path: "benchcmp/cold-index/chunks/first.bin".to_owned(),
            object_size: 64,
            ..Default::default()
        };
        let second = ColdChunkRef {
            start_offset: 64,
            end_offset: 128,
            s3_path: "benchcmp/cold-index/chunks/second.bin".to_owned(),
            object_size: 64,
            ..Default::default()
        };
        let replacement = ColdChunkRef {
            start_offset: 0,
            end_offset: 128,
            s3_path: "benchcmp/cold-index/chunks/compacted.bin".to_owned(),
            object_size: 128,
            ..Default::default()
        };
        for chunk in [&first, &second] {
            write_cold_chunk_index_pages(&store, &stream_id, chunk)
                .await
                .expect("write input chunk");
        }

        let rollback = replace_cold_chunk_index_pages_with_rollback(
            &store,
            &stream_id,
            &[first.clone(), second.clone()],
            &replacement,
        )
        .await
        .expect("replace chunks")
        .expect("inputs still match");
        rollback_cold_index_pages(&store, rollback)
            .await
            .expect("rollback replacement");

        let page = store
            .get_page(&ColdIndexPageKey {
                stream_id,
                generation: 0,
                page_id: 0,
            })
            .await
            .expect("get page")
            .expect("page exists");
        assert_eq!(page.cold_chunks, vec![first, second]);
    }

    #[tokio::test]
    async fn read_reload_repairs_stale_cached_page() {
        let store = Arc::new(InMemoryColdIndexPageStore::new());
        let stream_id = BucketStreamId::new("benchcmp", "cold-index");
        let cache = ColdIndexPageCache::new(store.clone(), 8);
        let first = ColdChunkRef {
            start_offset: 0,
            end_offset: 128,
            s3_path: "benchcmp/cold-index/chunks/first.bin".to_owned(),
            object_size: 128,
            ..Default::default()
        };
        write_cold_chunk_index_pages(store.as_ref(), &stream_id, &first)
            .await
            .expect("write first chunk");
        assert_eq!(
            cache
                .object_segments_for_read(&stream_id, &StreamReadColdIndexSegment {
                    generation: 0,
                    page_id: 0,
                    read_start_offset: 0,
                    len: 1,
                },)
                .await
                .expect("read first byte")
                .len(),
            1
        );

        let second = ColdChunkRef {
            start_offset: 128,
            end_offset: 256,
            s3_path: "benchcmp/cold-index/chunks/second.bin".to_owned(),
            object_size: 128,
            ..Default::default()
        };
        write_cold_chunk_index_pages(store.as_ref(), &stream_id, &second)
            .await
            .expect("write second chunk behind cache");

        let objects = cache
            .object_segments_for_read(&stream_id, &StreamReadColdIndexSegment {
                generation: 0,
                page_id: 0,
                read_start_offset: 128,
                len: 1,
            })
            .await
            .expect("reload stale page");
        assert_eq!(objects[0].s3_path, second.s3_path);
    }

    #[test]
    fn binary_page_format_round_trips_and_validates() {
        let key = key(42);
        let mut page = page(128, 256);
        page.external_segments.push(ObjectPayloadRef {
            start_offset: 256,
            end_offset: 300,
            s3_path: "benchcmp/cold-index/external/00000000000000000256.bin".to_owned(),
            object_size: 44,
            ..Default::default()
        });
        let bytes = encode_page(&key, &page);
        assert!(bytes.starts_with(COLD_INDEX_PAGE_MAGIC));

        assert_eq!(decode_page(&key, &bytes).expect("decode page"), page);

        let mut corrupted = bytes.clone();
        let last = corrupted.last_mut().expect("checksum byte");
        *last ^= 0xff;
        let err = decode_page(&key, &corrupted).expect_err("corrupt checksum");
        assert_eq!(err.kind(), io::ErrorKind::InvalidData);

        let wrong_key = ColdIndexPageKey {
            stream_id: key.stream_id.clone(),
            generation: key.generation + 1,
            page_id: key.page_id,
        };
        let err = decode_page(&wrong_key, &bytes).expect_err("key mismatch");
        assert_eq!(err.kind(), io::ErrorKind::InvalidData);
    }

    #[test]
    fn affinity_is_part_of_the_page_path_and_binary_identity() {
        let key = ColdIndexPageKey {
            stream_id: BucketStreamId::with_affinity("benchcmp", "run-42", "journal"),
            generation: 7,
            page_id: 42,
        };
        assert_eq!(
            key.path(),
            "benchcmp/run-42/journal/cold-index/00000000000000000007/00000000000000000042.idx"
        );

        let page = page(0, 10);
        let bytes = encode_page(&key, &page);
        assert_eq!(decode_page(&key, &bytes).expect("decode page"), page);

        let wrong_key = ColdIndexPageKey {
            stream_id: BucketStreamId::with_affinity("benchcmp", "run-43", "journal"),
            ..key
        };
        assert_eq!(
            decode_page(&wrong_key, &bytes)
                .expect_err("affinity mismatch")
                .kind(),
            io::ErrorKind::InvalidData
        );
    }

    #[tokio::test]
    async fn page_cache_loads_on_miss_and_evicts_lru() {
        let store = Arc::new(InMemoryColdIndexPageStore::new());
        for page_id in 0..3 {
            store
                .put_page(&key(page_id), &page(page_id * 100, page_id * 100 + 100))
                .await
                .expect("put page");
        }
        let cache = ColdIndexPageCache::new(store, 2);

        assert_eq!(
            cache
                .get_page(&key(0))
                .await
                .expect("load page")
                .expect("page")
                .start_offset,
            0
        );
        assert_eq!(
            cache
                .get_page(&key(1))
                .await
                .expect("load page")
                .expect("page")
                .start_offset,
            100
        );
        assert_eq!(cache.cached_page_count(), 2);

        // Touch page 0 so page 1 becomes the eviction candidate.
        assert!(
            cache
                .get_page(&key(0))
                .await
                .expect("cached page")
                .is_some()
        );
        assert_eq!(
            cache
                .get_page(&key(2))
                .await
                .expect("load page")
                .expect("page")
                .start_offset,
            200
        );
        assert_eq!(cache.cached_page_count(), 2);
    }

    #[tokio::test]
    async fn zero_capacity_cache_does_not_retain_pages() {
        let store = Arc::new(InMemoryColdIndexPageStore::new());
        store
            .put_page(&key(0), &page(0, 64))
            .await
            .expect("put page");
        let cache = ColdIndexPageCache::new(store, 0);

        assert!(cache.get_page(&key(0)).await.expect("load page").is_some());
        assert_eq!(cache.cached_page_count(), 0);
    }

    #[tokio::test]
    async fn selects_and_replaces_target_sized_contiguous_chunks() {
        let store = InMemoryColdIndexPageStore::new();
        let stream_id = BucketStreamId::new("benchcmp", "compact");
        let chunks = (0..4)
            .map(|index| ColdChunkRef {
                start_offset: index * 2,
                end_offset: index * 2 + 2,
                object_size: 2,
                s3_path: format!("old-{index}"),
                ..Default::default()
            })
            .collect::<Vec<_>>();
        for chunk in &chunks {
            write_cold_chunk_index_pages(&store, &stream_id, chunk)
                .await
                .expect("write chunk index");
        }
        let selected =
            select_cold_chunk_compaction(&chunks, 8, 16).expect("select compaction candidate");
        assert_eq!(selected, chunks);
        let replacement = ColdChunkRef {
            start_offset: 0,
            end_offset: 8,
            object_size: 8,
            s3_path: "replacement".to_owned(),
            ..Default::default()
        };
        assert!(
            replace_cold_chunk_index_pages(&store, &stream_id, &selected, &replacement)
                .await
                .expect("replace chunks")
        );
        let loaded = load_cold_chunks_from_pages(&store, &[ColdIndexPageKey {
            stream_id,
            generation: 0,
            page_id: 0,
        }])
        .await
        .expect("load replacement");
        assert_eq!(loaded, vec![replacement]);
    }
}