repo-stream / src / disk.rs

/*!
Disk storage for blocks on disk

Currently this uses sqlite. In testing sqlite wasn't the fastest, but it seemed
to be the best behaved in terms of both on-disk space usage and memory usage.

```no_run
# use repo_stream::{DiskBuilder, DiskError};
# #[tokio::main]
# async fn main() -> Result<(), DiskError> {
let store = DiskBuilder::new()
    .with_max_stored_mb(1024) // errors when >1GiB of processed blocks are inserted
    .open("/some/path.db".into()).await?;
# Ok(())
# }
```
*/

use crate::drive::DriveError;
use candystore::{CandyError, CandyStore, Config};
use std::path::PathBuf;

#[derive(Debug, thiserror::Error)]
pub enum DiskError {
    /// A wrapped database error
    ///
    /// (The wrapped err should probably be obscured to remove public-facing
    /// sqlite bits)
    #[error(transparent)]
    DbError(#[from] CandyError),
    /// Unfortunately candystore uses anyhow::Result for it's open call
    #[error("Failed on a db call, see logs")]
    DbGarbageError,
    /// A tokio blocking task failed to join
    #[error("Failed to join a tokio blocking task: {0}")]
    JoinError(#[from] tokio::task::JoinError),
    /// The total size of stored blocks exceeded the allowed size
    ///
    /// If you need to process *really* big CARs, you can configure a higher
    /// limit.
    #[error("Maximum disk size reached")]
    MaxSizeExceeded,
}

/// Builder-style disk store setup
#[derive(Debug, Clone)]
pub struct DiskBuilder {
    /// Database stored block size limit
    ///
    /// Default: 10 GiB
    ///
    /// Note: actual size on disk may be more, but should approximately scale
    /// with this limit
    pub max_stored_mb: usize,
}

impl Default for DiskBuilder {
    fn default() -> Self {
        Self {
            max_stored_mb: 10 * 1024, // 10 GiB
        }
    }
}

impl DiskBuilder {
    /// Begin configuring the storage with defaults
    pub fn new() -> Self {
        Default::default()
    }
    /// Set the approximate stored block size limit
    ///
    /// Default: 10 GiB
    pub fn with_max_stored_mb(mut self, max: usize) -> Self {
        self.max_stored_mb = max;
        self
    }
    /// Open and initialize the actual disk storage
    pub async fn open(&self, path: PathBuf, keys_hint: Option<usize>) -> Result<DiskStore, DiskError> {
        DiskStore::new(path, self.max_stored_mb, keys_hint).await
    }
}

/// On-disk block storage
pub struct DiskStore {
    db: CandyStore,
    max_stored: usize,
    stored: usize,
}

impl DiskStore {
    /// Initialize a new disk store
    pub async fn new(path: PathBuf, max_stored_mb: usize, keys_hint: Option<usize>) -> Result<Self, DiskError> {
        let max_stored = max_stored_mb * 2_usize.pow(20);
        let db = tokio::task::spawn_blocking(move || {
            let mut conf = Config::default();
            // conf.max_shard_size = 256 * 1024 * 1024;
            // conf.min_compaction_threashold = 32 * 1024 * 1024;
            // conf.expected_number_of_keys = 1_200_000;
            if let Some(hint) = keys_hint {
                conf.expected_number_of_keys = hint;
            }
            conf.num_compaction_threads = 1;
            let db = CandyStore::open(path, conf).map_err(|e| {
                log::error!("{e:?}");
                DiskError::DbGarbageError
            })?;

            Ok::<_, DiskError>(db)
        })
        .await??;

        Ok(Self {
            db,
            max_stored,
            stored: 0,
        })
    }

    /// Drop and recreate the kv table
    pub async fn reset(self) -> Result<Self, DiskError> {
        tokio::task::spawn_blocking(move || {
            Self::reset_tables(&self.db)?;
            Ok(self)
        })
        .await?
    }
    fn reset_tables(db: &CandyStore) -> Result<(), DiskError> {
        db.clear().map_err(|e| {
            log::error!("{e:?}");
            DiskError::DbGarbageError
        })?;

        Ok(())
    }

    pub(crate) fn put_many(
        &mut self,
        kv: impl Iterator<Item = Result<(Vec<u8>, Vec<u8>), DriveError>>,
    ) -> Result<(), DriveError> {
        for pair in kv {
            let (k, v) = pair?;
            self.stored += v.len();
            if self.stored > self.max_stored {
                return Err(DiskError::MaxSizeExceeded.into());
            }
            self.db.owned_set(k, &v).map_err(|e| {
                log::error!("{e:?}");
                DiskError::DbGarbageError
            })?;
        }
        Ok(())
    }
    pub(crate) fn get(&mut self, key: Vec<u8>) -> Result<Option<Vec<u8>>, DiskError> {
        self.db.owned_get(key).map_err(|e| {
            log::error!("{e:?}");
            DiskError::DbGarbageError
        })
    }
}