redb-turbo 0.2.0

use crate::page_crypto::{PageCompression, PageCrypto};
use crate::tree_store::page_store::base::PageHint;
use crate::tree_store::page_store::lru_cache::LRUCache;
use crate::{CacheStats, DatabaseError, Result, StorageBackend, StorageError};
use std::ops::{Index, IndexMut};
use std::slice::SliceIndex;
#[cfg(feature = "cache_metrics")]
use std::sync::atomic::AtomicU64;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, RwLock};

pub(super) struct WritablePage {
    buffer: Arc<Mutex<LRUWriteCache>>,
    offset: u64,
    data: Arc<[u8]>,
}

impl WritablePage {
    pub(super) fn mem(&self) -> &[u8] {
        &self.data
    }

    pub(super) fn mem_mut(&mut self) -> &mut [u8] {
        Arc::get_mut(&mut self.data).unwrap()
    }
}

impl Drop for WritablePage {
    fn drop(&mut self) {
        self.buffer
            .lock()
            .unwrap()
            .return_value(self.offset, self.data.clone());
    }
}

impl<I: SliceIndex<[u8]>> Index<I> for WritablePage {
    type Output = I::Output;

    fn index(&self, index: I) -> &Self::Output {
        self.mem().index(index)
    }
}

impl<I: SliceIndex<[u8]>> IndexMut<I> for WritablePage {
    fn index_mut(&mut self, index: I) -> &mut Self::Output {
        self.mem_mut().index_mut(index)
    }
}

#[derive(Default)]
struct LRUWriteCache {
    cache: LRUCache<Option<Arc<[u8]>>>,
}

impl LRUWriteCache {
    fn new() -> Self {
        Self {
            cache: Default::default(),
        }
    }

    fn insert(&mut self, key: u64, value: Arc<[u8]>) {
        assert!(self.cache.insert(key, Some(value)).is_none());
    }

    fn get(&self, key: u64) -> Option<&Arc<[u8]>> {
        self.cache.get(key).map(|x| x.as_ref().unwrap())
    }

    fn remove(&mut self, key: u64) -> Option<Arc<[u8]>> {
        if let Some(value) = self.cache.remove(key) {
            assert!(value.is_some());
            return value;
        }
        None
    }

    fn return_value(&mut self, key: u64, value: Arc<[u8]>) {
        assert!(self.cache.get_mut(key).unwrap().replace(value).is_none());
    }

    fn take_value(&mut self, key: u64) -> Option<Arc<[u8]>> {
        if let Some(value) = self.cache.get_mut(key) {
            let result = value.take().unwrap();
            return Some(result);
        }
        None
    }

    fn pop_lowest_priority(&mut self) -> Option<(u64, Arc<[u8]>)> {
        for _ in 0..self.cache.len() {
            if let Some((k, v)) = self.cache.pop_lowest_priority() {
                if let Some(v_inner) = v {
                    return Some((k, v_inner));
                }

                // Value is borrowed by take_value(). We can't evict it, so put it back.
                self.cache.insert(k, v);
            } else {
                break;
            }
        }
        None
    }

    fn clear(&mut self) {
        self.cache.clear();
    }
}

struct CheckedBackend {
    file: Box<dyn StorageBackend>,
    io_failed: AtomicBool,
    crypto: Option<Arc<dyn PageCrypto>>,
    compression: Option<Arc<dyn PageCompression>>,
    page_size: usize,
}

impl std::fmt::Debug for CheckedBackend {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CheckedBackend")
            .field("io_failed", &self.io_failed)
            .field("has_crypto", &self.crypto.is_some())
            .field("has_compression", &self.compression.is_some())
            .field("page_size", &self.page_size)
            .finish()
    }
}

impl CheckedBackend {
    fn new(file: Box<dyn StorageBackend>, page_size: usize) -> Self {
        Self {
            file,
            io_failed: AtomicBool::new(false),
            crypto: None,
            compression: None,
            page_size,
        }
    }

    fn with_transforms(
        file: Box<dyn StorageBackend>,
        page_size: usize,
        crypto: Option<Arc<dyn PageCrypto>>,
        compression: Option<Arc<dyn PageCompression>>,
    ) -> Self {
        Self {
            file,
            io_failed: AtomicBool::new(false),
            crypto,
            compression,
            page_size,
        }
    }

    fn set_failure(&self) {
        self.io_failed.store(true, Ordering::Release);
    }

    fn check_failure(&self) -> Result<()> {
        if self.io_failed.load(Ordering::Acquire) {
            Err(StorageError::PreviousIo)
        } else {
            Ok(())
        }
    }

    fn len(&self) -> Result<u64> {
        self.check_failure()?;
        let result = self.file.len();
        if result.is_err() {
            self.io_failed.store(true, Ordering::Release);
        }
        result.map_err(StorageError::from)
    }

    fn read(&self, offset: u64, len: usize) -> Result<Vec<u8>> {
        self.check_failure()?;
        let result = self.file.read(offset, len);
        if result.is_err() {
            self.io_failed.store(true, Ordering::Release);
        }
        let raw_data = result.map_err(StorageError::from)?;

        // Apply transforms in reverse order: decrypt first, then decompress
        // (write order is: compress -> encrypt -> disk)

        if len == self.page_size {
            // Single page read - apply both decryption and decompression
            let mut data = raw_data;

            // Decrypt if crypto is configured
            if let Some(ref crypto) = self.crypto {
                if offset >= crypto.encryption_start_offset() {
                    data = crypto
                        .decrypt(offset, &data, self.page_size)
                        .map_err(StorageError::Io)?;
                }
            }

            // Decompress if compression is configured
            if let Some(ref compression) = self.compression {
                if offset >= compression.compression_start_offset() {
                    data = compression
                        .decompress(offset, &data, self.page_size)
                        .map_err(StorageError::Io)?;
                }
            }

            Ok(data)
        } else if len % self.page_size == 0 && (self.compression.is_some() || self.crypto.is_some()) {
            // Multi-page read with compression/encryption - process page-by-page
            let num_pages = len / self.page_size;
            let mut result = Vec::with_capacity(len);

            for i in 0..num_pages {
                let page_offset = offset + (i * self.page_size) as u64;
                let page_start = i * self.page_size;
                let mut page_data = raw_data[page_start..page_start + self.page_size].to_vec();

                // Decrypt first (reverse of write order: compress -> encrypt)
                if let Some(ref crypto) = self.crypto {
                    if page_offset >= crypto.encryption_start_offset() {
                        page_data = crypto
                            .decrypt(page_offset, &page_data, self.page_size)
                            .map_err(StorageError::Io)?;
                    }
                }

                // Then decompress
                if let Some(ref compression) = self.compression {
                    if page_offset >= compression.compression_start_offset() {
                        page_data = compression
                            .decompress(page_offset, &page_data, self.page_size)
                            .map_err(StorageError::Io)?;
                    }
                }

                result.extend_from_slice(&page_data);
            }

            Ok(result)
        } else {
            // Non-page-aligned read - return as-is
            Ok(raw_data)
        }
    }

    fn set_len(&self, len: u64) -> Result<()> {
        self.check_failure()?;
        let result = self.file.set_len(len);
        if result.is_err() {
            self.io_failed.store(true, Ordering::Release);
        }
        result.map_err(StorageError::from)
    }

    fn sync_data(&self, eventual: bool) -> Result<()> {
        self.check_failure()?;
        let result = self.file.sync_data(eventual);
        if result.is_err() {
            self.io_failed.store(true, Ordering::Release);
        }
        result.map_err(StorageError::from)
    }

    fn write(&self, offset: u64, data: &[u8]) -> Result<()> {
        self.check_failure()?;

        // Apply transforms: compress first, then encrypt
        let mut data_to_write: std::borrow::Cow<'_, [u8]> = std::borrow::Cow::Borrowed(data);

        // Compress if compression is configured and this is a full page write
        if let Some(ref compression) = self.compression {
            if data.len() == self.page_size && offset >= compression.compression_start_offset() {
                let compressed = compression
                    .compress(offset, &data_to_write, self.page_size)
                    .map_err(StorageError::Io)?;
                data_to_write = std::borrow::Cow::Owned(compressed);
            }
        }

        // Encrypt if crypto is configured and this is a full page write
        if let Some(ref crypto) = self.crypto {
            if data_to_write.len() == self.page_size && offset >= crypto.encryption_start_offset() {
                let encrypted = crypto
                    .encrypt(offset, &data_to_write, self.page_size)
                    .map_err(StorageError::Io)?;
                data_to_write = std::borrow::Cow::Owned(encrypted);
            }
        }

        let result = self.file.write(offset, &data_to_write);
        if result.is_err() {
            self.io_failed.store(true, Ordering::Release);
        }
        result.map_err(StorageError::from)
    }

    fn write_batch(&self, ops: &[(u64, Arc<[u8]>)]) -> Result<()> {
        self.check_failure()?;

        // Apply transforms: compress first, then encrypt
        // For multi-page writes, we process compression page-by-page but NOT encryption.
        // Encryption has per-page overhead that requires data to be stored accounting for it,
        // which is handled by redb at a higher level for single-page writes.
        let mut transformed_ops: Vec<(u64, Vec<u8>)> = Vec::new();

        for (offset, data) in ops.iter() {
            if data.len() == self.page_size {
                // Single page write - process directly (both compression and encryption)
                let mut transformed: Vec<u8> = data.to_vec();

                // Compress if configured
                if let Some(ref compression) = self.compression {
                    if *offset >= compression.compression_start_offset() {
                        transformed = compression
                            .compress(*offset, &transformed, self.page_size)
                            .map_err(StorageError::Io)?;
                    }
                }

                // Encrypt if configured (only for single pages)
                if let Some(ref crypto) = self.crypto {
                    if transformed.len() == self.page_size && *offset >= crypto.encryption_start_offset() {
                        transformed = crypto
                            .encrypt(*offset, &transformed, self.page_size)
                            .map_err(StorageError::Io)?;
                    }
                }

                transformed_ops.push((*offset, transformed));
            } else if data.len() % self.page_size == 0 && (self.compression.is_some() || self.crypto.is_some()) {
                // Multi-page write with compression/encryption - process page-by-page
                let num_pages = data.len() / self.page_size;
                for i in 0..num_pages {
                    let page_offset = *offset + (i * self.page_size) as u64;
                    let page_start = i * self.page_size;
                    let page_data = &data[page_start..page_start + self.page_size];
                    let mut transformed: Vec<u8> = page_data.to_vec();

                    // Compress first
                    if let Some(ref compression) = self.compression {
                        if page_offset >= compression.compression_start_offset() {
                            transformed = compression
                                .compress(page_offset, &transformed, self.page_size)
                                .map_err(StorageError::Io)?;
                        }
                    }

                    // Then encrypt
                    if let Some(ref crypto) = self.crypto {
                        if transformed.len() == self.page_size && page_offset >= crypto.encryption_start_offset() {
                            transformed = crypto
                                .encrypt(page_offset, &transformed, self.page_size)
                                .map_err(StorageError::Io)?;
                        }
                    }

                    transformed_ops.push((page_offset, transformed));
                }
            } else {
                // Non-page-aligned write - pass through as-is
                transformed_ops.push((*offset, data.to_vec()));
            }
        }

        // Convert to the format expected by StorageBackend::write_batch
        let batch_refs: Vec<(u64, &[u8])> = transformed_ops
            .iter()
            .map(|(offset, data)| (*offset, data.as_slice()))
            .collect();

        let result = self.file.write_batch(&batch_refs);
        if result.is_err() {
            self.io_failed.store(true, Ordering::Release);
        }
        result.map_err(StorageError::from)
    }
}

pub(super) struct PagedCachedFile {
    file: CheckedBackend,
    page_size: u64,
    max_read_cache_bytes: usize,
    read_cache_bytes: AtomicUsize,
    max_write_buffer_bytes: usize,
    write_buffer_bytes: AtomicUsize,
    #[cfg(feature = "cache_metrics")]
    reads_total: AtomicU64,
    #[cfg(feature = "cache_metrics")]
    reads_hits: AtomicU64,
    #[cfg(feature = "cache_metrics")]
    evictions: AtomicU64,
    read_cache: Vec<RwLock<LRUCache<Arc<[u8]>>>>,
    // TODO: maybe move this cache to WriteTransaction?
    write_buffer: Arc<Mutex<LRUWriteCache>>,
}

impl PagedCachedFile {
    #[allow(dead_code)] // Used in tests
    pub(super) fn new(
        file: Box<dyn StorageBackend>,
        page_size: u64,
        max_read_cache_bytes: usize,
        max_write_buffer_bytes: usize,
    ) -> Result<Self, DatabaseError> {
        Self::new_with_transforms(file, page_size, max_read_cache_bytes, max_write_buffer_bytes, None, None)
    }

    pub(super) fn new_with_transforms(
        file: Box<dyn StorageBackend>,
        page_size: u64,
        max_read_cache_bytes: usize,
        max_write_buffer_bytes: usize,
        crypto: Option<Arc<dyn PageCrypto>>,
        compression: Option<Arc<dyn PageCompression>>,
    ) -> Result<Self, DatabaseError> {
        let read_cache = (0..Self::lock_stripes())
            .map(|_| RwLock::new(LRUCache::new()))
            .collect();

        let checked_backend = if crypto.is_some() || compression.is_some() {
            CheckedBackend::with_transforms(file, page_size as usize, crypto, compression)
        } else {
            CheckedBackend::new(file, page_size as usize)
        };

        Ok(Self {
            file: checked_backend,
            page_size,
            max_read_cache_bytes,
            read_cache_bytes: AtomicUsize::new(0),
            max_write_buffer_bytes,
            write_buffer_bytes: AtomicUsize::new(0),
            #[cfg(feature = "cache_metrics")]
            reads_total: Default::default(),
            #[cfg(feature = "cache_metrics")]
            reads_hits: Default::default(),
            #[cfg(feature = "cache_metrics")]
            evictions: Default::default(),
            read_cache,
            write_buffer: Arc::new(Mutex::new(LRUWriteCache::new())),
        })
    }

    #[allow(clippy::unused_self)]
    pub(crate) fn cache_stats(&self) -> CacheStats {
        CacheStats {
            #[cfg(not(feature = "cache_metrics"))]
            evictions: 0,
            #[cfg(feature = "cache_metrics")]
            evictions: self.evictions.load(Ordering::Acquire),
        }
    }

    pub(crate) fn check_io_errors(&self) -> Result {
        self.file.check_failure()
    }

    pub(crate) fn set_irrecoverable_io_error(&self) {
        self.file.set_failure();
    }

    pub(crate) fn raw_file_len(&self) -> Result<u64> {
        self.file.len()
    }

    const fn lock_stripes() -> u64 {
        131
    }

    fn flush_write_buffer(&self) -> Result {
        let mut write_buffer = self.write_buffer.lock().unwrap();

        // Collect all writes for batch operation
        let batch_ops: Vec<(u64, Arc<[u8]>)> = write_buffer
            .cache
            .iter()
            .map(|(offset, buffer)| (*offset, buffer.as_ref().unwrap().clone()))
            .collect();

        // Batch write all pages (uses io_uring on Linux if available)
        self.file.write_batch(&batch_ops)?;

        for (offset, buffer) in write_buffer.cache.iter_mut() {
            let buffer = buffer.take().unwrap();
            let cache_size = self
                .read_cache_bytes
                .fetch_add(buffer.len(), Ordering::AcqRel);

            if cache_size + buffer.len() <= self.max_read_cache_bytes {
                let cache_slot: usize = (offset % Self::lock_stripes()).try_into().unwrap();
                let mut lock = self.read_cache[cache_slot].write().unwrap();
                if let Some(replaced) = lock.insert(*offset, buffer) {
                    // A race could cause us to replace an existing buffer
                    self.read_cache_bytes
                        .fetch_sub(replaced.len(), Ordering::AcqRel);
                }
            } else {
                self.read_cache_bytes
                    .fetch_sub(buffer.len(), Ordering::AcqRel);
                break;
            }
        }
        self.write_buffer_bytes.store(0, Ordering::Release);
        write_buffer.clear();

        Ok(())
    }

    // Caller should invalidate all cached pages that are no longer valid
    pub(super) fn resize(&self, len: u64) -> Result {
        // TODO: be more fine-grained about this invalidation
        self.invalidate_cache_all();

        self.file.set_len(len)
    }

    pub(super) fn flush(&self, #[allow(unused_variables)] eventual: bool) -> Result {
        self.flush_write_buffer()?;

        self.file.sync_data(eventual)
    }

    // Make writes visible to readers, but does not guarantee any durability
    pub(super) fn write_barrier(&self) -> Result {
        self.flush_write_buffer()
    }

    // Read directly from the file, ignoring any cached data
    pub(super) fn read_direct(&self, offset: u64, len: usize) -> Result<Vec<u8>> {
        self.file.read(offset, len)
    }

    // Read with caching. Caller must not read overlapping ranges without first calling invalidate_cache().
    // Doing so will not cause UB, but is a logic error.
    pub(super) fn read(&self, offset: u64, len: usize, hint: PageHint) -> Result<Arc<[u8]>> {
        debug_assert_eq!(0, offset % self.page_size);
        #[cfg(feature = "cache_metrics")]
        self.reads_total.fetch_add(1, Ordering::AcqRel);

        if !matches!(hint, PageHint::Clean) {
            let lock = self.write_buffer.lock().unwrap();
            if let Some(cached) = lock.get(offset) {
                #[cfg(feature = "cache_metrics")]
                self.reads_hits.fetch_add(1, Ordering::Release);
                debug_assert_eq!(cached.len(), len);
                return Ok(cached.clone());
            }
        }

        let cache_slot: usize = (offset % Self::lock_stripes()).try_into().unwrap();
        {
            let read_lock = self.read_cache[cache_slot].read().unwrap();
            if let Some(cached) = read_lock.get(offset) {
                #[cfg(feature = "cache_metrics")]
                self.reads_hits.fetch_add(1, Ordering::Release);
                debug_assert_eq!(cached.len(), len);
                return Ok(cached.clone());
            }
        }

        let buffer: Arc<[u8]> = self.read_direct(offset, len)?.into();
        let cache_size = self.read_cache_bytes.fetch_add(len, Ordering::AcqRel);
        let mut write_lock = self.read_cache[cache_slot].write().unwrap();
        let cache_size = if let Some(replaced) = write_lock.insert(offset, buffer.clone()) {
            // A race could cause us to replace an existing buffer
            self.read_cache_bytes
                .fetch_sub(replaced.len(), Ordering::AcqRel)
        } else {
            cache_size
        };
        let mut removed = 0;
        if cache_size + len > self.max_read_cache_bytes {
            while removed < len {
                if let Some((_, v)) = write_lock.pop_lowest_priority() {
                    #[cfg(feature = "cache_metrics")]
                    {
                        self.evictions.fetch_add(1, Ordering::Relaxed);
                    }
                    removed += v.len();
                } else {
                    break;
                }
            }
        }
        if removed > 0 {
            self.read_cache_bytes.fetch_sub(removed, Ordering::AcqRel);
        }

        Ok(buffer)
    }

    // Discard pending writes to the given range
    pub(super) fn cancel_pending_write(&self, offset: u64, _len: usize) {
        assert_eq!(0, offset % self.page_size);
        if let Some(removed) = self.write_buffer.lock().unwrap().remove(offset) {
            self.write_buffer_bytes
                .fetch_sub(removed.len(), Ordering::Release);
        }
    }

    // Invalidate any caching of the given range. After this call overlapping reads of the range are allowed
    //
    // NOTE: Invalidating a cached region in subsections is permitted, as long as all subsections are invalidated
    pub(super) fn invalidate_cache(&self, offset: u64, len: usize) {
        let cache_slot: usize = (offset % Self::lock_stripes()).try_into().unwrap();
        let mut lock = self.read_cache[cache_slot].write().unwrap();
        if let Some(removed) = lock.remove(offset) {
            assert_eq!(len, removed.len());
            self.read_cache_bytes
                .fetch_sub(removed.len(), Ordering::AcqRel);
        }
    }

    pub(super) fn invalidate_cache_all(&self) {
        for cache_slot in 0..self.read_cache.len() {
            let mut lock = self.read_cache[cache_slot].write().unwrap();
            while let Some((_, removed)) = lock.pop_lowest_priority() {
                self.read_cache_bytes
                    .fetch_sub(removed.len(), Ordering::AcqRel);
            }
        }
    }

    // If overwrite is true, the page is initialized to zero
    // cache_policy takes the existing data as an argument and returns the priority. The priority should be stable and not change after WritablePage is dropped
    pub(super) fn write(&self, offset: u64, len: usize, overwrite: bool) -> Result<WritablePage> {
        assert_eq!(0, offset % self.page_size);
        let mut lock = self.write_buffer.lock().unwrap();

        // TODO: allow hint that page is known to be dirty and will not be in the read cache
        let cache_slot: usize = (offset % Self::lock_stripes()).try_into().unwrap();
        let existing = {
            let mut lock = self.read_cache[cache_slot].write().unwrap();
            if let Some(removed) = lock.remove(offset) {
                assert_eq!(
                    len,
                    removed.len(),
                    "cache inconsistency {len} != {} for offset {offset}",
                    removed.len()
                );
                self.read_cache_bytes
                    .fetch_sub(removed.len(), Ordering::AcqRel);
                Some(removed)
            } else {
                None
            }
        };

        let data = if let Some(removed) = lock.take_value(offset) {
            removed
        } else {
            let previous = self.write_buffer_bytes.fetch_add(len, Ordering::AcqRel);
            if previous + len > self.max_write_buffer_bytes {
                let mut removed_bytes = 0;
                while removed_bytes < len {
                    if let Some((offset, buffer)) = lock.pop_lowest_priority() {
                        let removed_len = buffer.len();
                        let result = self.file.write(offset, &buffer);
                        if result.is_err() {
                            lock.insert(offset, buffer);
                        }
                        result?;
                        self.write_buffer_bytes
                            .fetch_sub(removed_len, Ordering::Release);
                        #[cfg(feature = "cache_metrics")]
                        {
                            self.evictions.fetch_add(1, Ordering::Relaxed);
                        }
                        removed_bytes += removed_len;
                    } else {
                        break;
                    }
                }
            }
            let result = if let Some(data) = existing {
                data
            } else if overwrite {
                vec![0; len].into()
            } else {
                self.read_direct(offset, len)?.into()
            };
            lock.insert(offset, result);
            lock.take_value(offset).unwrap()
        };
        Ok(WritablePage {
            buffer: self.write_buffer.clone(),
            offset,
            data,
        })
    }
}

#[cfg(test)]
mod test {
    use crate::StorageBackend;
    use crate::backends::InMemoryBackend;
    use crate::tree_store::PageHint;
    use crate::tree_store::page_store::cached_file::PagedCachedFile;
    use std::sync::Arc;
    use std::sync::atomic::Ordering;

    #[test]
    fn cache_leak() {
        let backend = InMemoryBackend::new();
        backend.set_len(1024).unwrap();
        let cached_file = PagedCachedFile::new(Box::new(backend), 128, 1024, 128).unwrap();
        let cached_file = Arc::new(cached_file);

        let t1 = {
            let cached_file = cached_file.clone();
            std::thread::spawn(move || {
                for _ in 0..1000 {
                    cached_file.read(0, 128, PageHint::None).unwrap();
                    cached_file.invalidate_cache(0, 128);
                }
            })
        };
        let t2 = {
            let cached_file = cached_file.clone();
            std::thread::spawn(move || {
                for _ in 0..1000 {
                    cached_file.read(0, 128, PageHint::None).unwrap();
                    cached_file.invalidate_cache(0, 128);
                }
            })
        };

        t1.join().unwrap();
        t2.join().unwrap();
        cached_file.invalidate_cache(0, 128);
        assert_eq!(cached_file.read_cache_bytes.load(Ordering::Acquire), 0);
    }
}