tranquil-pds/crates/tranquil-repo/src/lib.rs

use bytes::Bytes;
use cid::Cid;
use jacquard_repo::error::RepoError;
use jacquard_repo::repo::CommitData;
use jacquard_repo::storage::BlockStore;
use multihash::Multihash;
use sha2::{Digest, Sha256};
use sqlx::PgPool;
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};

#[derive(Clone)]
pub struct PostgresBlockStore {
    pool: PgPool,
}

impl PostgresBlockStore {
    pub fn new(pool: PgPool) -> Self {
        Self { pool }
    }

    pub fn pool(&self) -> &PgPool {
        &self.pool
    }
}

impl BlockStore for PostgresBlockStore {
    async fn get(&self, cid: &Cid) -> Result<Option<Bytes>, RepoError> {
        let cid_bytes = cid.to_bytes();
        let row = sqlx::query!("SELECT data FROM blocks WHERE cid = $1", &cid_bytes)
            .fetch_optional(&self.pool)
            .await
            .map_err(RepoError::storage)?;
        match row {
            Some(row) => Ok(Some(Bytes::from(row.data))),
            None => Ok(None),
        }
    }

    async fn put(&self, data: &[u8]) -> Result<Cid, RepoError> {
        let mut hasher = Sha256::new();
        hasher.update(data);
        let hash = hasher.finalize();
        let multihash = Multihash::wrap(0x12, &hash).map_err(|e| {
            RepoError::storage(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("Failed to wrap multihash: {:?}", e),
            ))
        })?;
        let cid = Cid::new_v1(0x71, multihash);
        let cid_bytes = cid.to_bytes();
        sqlx::query!(
            "INSERT INTO blocks (cid, data) VALUES ($1, $2) ON CONFLICT (cid) DO NOTHING",
            &cid_bytes,
            data
        )
        .execute(&self.pool)
        .await
        .map_err(RepoError::storage)?;
        Ok(cid)
    }

    async fn has(&self, cid: &Cid) -> Result<bool, RepoError> {
        let cid_bytes = cid.to_bytes();
        let row = sqlx::query!("SELECT 1 as one FROM blocks WHERE cid = $1", &cid_bytes)
            .fetch_optional(&self.pool)
            .await
            .map_err(RepoError::storage)?;
        Ok(row.is_some())
    }

    async fn put_many(
        &self,
        blocks: impl IntoIterator<Item = (Cid, Bytes)> + Send,
    ) -> Result<(), RepoError> {
        let blocks: Vec<_> = blocks.into_iter().collect();
        if blocks.is_empty() {
            return Ok(());
        }
        let cids: Vec<Vec<u8>> = blocks.iter().map(|(cid, _)| cid.to_bytes()).collect();
        let data: Vec<&[u8]> = blocks.iter().map(|(_, d)| d.as_ref()).collect();
        sqlx::query!(
            r#"
            INSERT INTO blocks (cid, data)
            SELECT * FROM UNNEST($1::bytea[], $2::bytea[])
            ON CONFLICT (cid) DO NOTHING
            "#,
            &cids,
            &data as &[&[u8]]
        )
        .execute(&self.pool)
        .await
        .map_err(RepoError::storage)?;
        Ok(())
    }

    async fn get_many(&self, cids: &[Cid]) -> Result<Vec<Option<Bytes>>, RepoError> {
        if cids.is_empty() {
            return Ok(Vec::new());
        }
        let cid_bytes: Vec<Vec<u8>> = cids.iter().map(|c| c.to_bytes()).collect();
        let rows = sqlx::query!(
            "SELECT cid, data FROM blocks WHERE cid = ANY($1)",
            &cid_bytes
        )
        .fetch_all(&self.pool)
        .await
        .map_err(RepoError::storage)?;
        let found: std::collections::HashMap<Vec<u8>, Bytes> = rows
            .into_iter()
            .map(|row| (row.cid, Bytes::from(row.data)))
            .collect();
        let results = cid_bytes
            .iter()
            .map(|cid| found.get(cid).cloned())
            .collect();
        Ok(results)
    }

    async fn apply_commit(&self, commit: CommitData) -> Result<(), RepoError> {
        self.put_many(commit.blocks).await?;
        Ok(())
    }
}

#[derive(Clone)]
pub struct TrackingBlockStore<S: BlockStore> {
    inner: S,
    written_blocks: Arc<Mutex<HashMap<Cid, Bytes>>>,
    read_cids: Arc<Mutex<HashSet<Cid>>>,
}

impl<S: BlockStore + Sync> TrackingBlockStore<S> {
    pub fn new(store: S) -> Self {
        Self {
            inner: store,
            written_blocks: Arc::new(Mutex::new(HashMap::new())),
            read_cids: Arc::new(Mutex::new(HashSet::new())),
        }
    }

    pub fn get_written_cids(&self) -> Vec<Cid> {
        match self.written_blocks.lock() {
            Ok(guard) => guard.keys().copied().collect(),
            Err(poisoned) => poisoned.into_inner().keys().copied().collect(),
        }
    }

    pub fn take_written_blocks(&self) -> HashMap<Cid, Bytes> {
        let mut guard = match self.written_blocks.lock() {
            Ok(g) => g,
            Err(poisoned) => poisoned.into_inner(),
        };
        std::mem::take(&mut *guard)
    }

    pub fn get_read_cids(&self) -> Vec<Cid> {
        match self.read_cids.lock() {
            Ok(guard) => guard.iter().cloned().collect(),
            Err(poisoned) => poisoned.into_inner().iter().cloned().collect(),
        }
    }

    pub fn get_all_relevant_cids(&self) -> Vec<Cid> {
        let written = self.get_written_cids();
        let read = self.get_read_cids();
        let mut all: HashSet<Cid> = written.into_iter().collect();
        all.extend(read);
        all.into_iter().collect()
    }
}

impl<S: BlockStore + Sync> BlockStore for TrackingBlockStore<S> {
    async fn get(&self, cid: &Cid) -> Result<Option<Bytes>, RepoError> {
        let result = self.inner.get(cid).await?;
        if result.is_some() {
            match self.read_cids.lock() {
                Ok(mut guard) => {
                    guard.insert(*cid);
                }
                Err(poisoned) => {
                    poisoned.into_inner().insert(*cid);
                }
            }
        }
        Ok(result)
    }

    async fn put(&self, data: &[u8]) -> Result<Cid, RepoError> {
        let cid = self.inner.put(data).await?;
        let bytes = Bytes::copy_from_slice(data);
        match self.written_blocks.lock() {
            Ok(mut guard) => {
                guard.insert(cid, bytes);
            }
            Err(poisoned) => {
                poisoned.into_inner().insert(cid, bytes);
            }
        }
        Ok(cid)
    }

    async fn has(&self, cid: &Cid) -> Result<bool, RepoError> {
        self.inner.has(cid).await
    }

    async fn put_many(
        &self,
        blocks: impl IntoIterator<Item = (Cid, Bytes)> + Send,
    ) -> Result<(), RepoError> {
        let blocks: Vec<_> = blocks.into_iter().collect();
        self.inner.put_many(blocks.clone()).await?;
        match self.written_blocks.lock() {
            Ok(mut guard) => guard.extend(blocks),
            Err(poisoned) => poisoned.into_inner().extend(blocks),
        }
        Ok(())
    }

    async fn get_many(&self, cids: &[Cid]) -> Result<Vec<Option<Bytes>>, RepoError> {
        let results = self.inner.get_many(cids).await?;
        cids.iter()
            .zip(results.iter())
            .filter(|(_, result)| result.is_some())
            .for_each(|(cid, _)| match self.read_cids.lock() {
                Ok(mut guard) => {
                    guard.insert(*cid);
                }
                Err(poisoned) => {
                    poisoned.into_inner().insert(*cid);
                }
            });
        Ok(results)
    }

    async fn apply_commit(&self, commit: CommitData) -> Result<(), RepoError> {
        self.put_many(commit.blocks).await?;
        Ok(())
    }
}