ordinary_storage/stores/

assets.rs

// Copyright (C) 2026 The Ordinary Authors.
//
// SPDX-License-Identifier: BSD-3-Clause

use anyhow::bail;
use bytes::{BufMut, Bytes, BytesMut};
use hashbrown::HashSet;
use ordinary_config::{AssetsLimits, CompressionAlgorithm};
use parking_lot::Mutex;
use saferlmdb::{
    self as lmdb, Database, DatabaseOptions, Environment, ReadTransaction, WriteTransaction, put,
};
use std::fmt::{Display, Formatter};
use std::sync::Arc;
use tracing::{Level, instrument};

pub struct PercentageDisplay(pub f64);

impl Display for PercentageDisplay {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:.2}%", self.0)
    }
}

pub struct AssetStore {
    limits: AssetsLimits,

    env: Arc<Environment>,

    /// stores static assets to be served by web server
    assets_db: Arc<Database<'static>>,

    log_sizes: bool,

    store_size: Arc<Mutex<u64>>,
}

impl AssetStore {
    pub fn new(
        limits: AssetsLimits,
        env: &Arc<Environment>,
        log_sizes: bool,
    ) -> anyhow::Result<Self> {
        let asset_db = Arc::new(Database::open(
            env.clone(),
            Some("asset"),
            &DatabaseOptions::new(lmdb::db::Flags::CREATE),
        )?);

        let mut store_size = 0;

        let txn = ReadTransaction::new(env.clone())?;
        let access = txn.access();

        let mut asset_cursor = txn.cursor(asset_db.clone())?;

        if let Ok((key, val)) = asset_cursor.first::<[u8], [u8]>(&access) {
            store_size += key.len() as u64;
            store_size += val.len() as u64;

            while let Ok((key, val)) = asset_cursor.next::<[u8], [u8]>(&access) {
                store_size += key.len() as u64;
                store_size += val.len() as u64;
            }
        }

        Ok(Self {
            limits,
            env: env.clone(),
            assets_db: asset_db,
            log_sizes,
            // todo: read in on start up
            store_size: Arc::new(Mutex::new(store_size)),
        })
    }

    #[allow(clippy::cast_precision_loss)]
    #[instrument(skip_all, err)]
    pub fn put(
        &self,
        path: &str,
        asset: &[u8],
        compression: Option<(&CompressionAlgorithm, usize)>,
    ) -> anyhow::Result<()> {
        let mut key = BytesMut::with_capacity(path.len() + 1);
        key.put(path.as_bytes());

        if let Some(compression) = compression {
            key.put_u8(compression.0.as_u8());
        }

        let size = (key.len() + asset.len()) as u64;

        if size > self.limits.max_asset_size {
            bail!("exceeds asset size limit");
        }

        let mut store_size = self.store_size.lock();

        if *store_size + size > self.limits.max_store_size {
            bail!("exceeds store size limit");
        }

        let mut overwrite_size = 0;

        let txn = WriteTransaction::new(self.env.clone())?;

        {
            let mut access = txn.access();
            if let Ok(result) = access.get::<[u8], [u8]>(&self.assets_db, key.as_ref()) {
                overwrite_size = (key.len() + result.len()) as u64;
            }

            access.put(&self.assets_db, key.as_ref(), asset, &put::Flags::empty())?;
        }

        txn.commit()?;

        *store_size -= overwrite_size;
        *store_size += size;

        if self.log_sizes {
            let (reduction, source_size) = if let Some((_, pre_compression_size)) = compression {
                (
                    Some(tracing::field::display(PercentageDisplay(
                        ((pre_compression_size as f64 - size as f64) / pre_compression_size as f64)
                            * 100.0,
                    ))),
                    pre_compression_size as u64,
                )
            } else {
                (None, size)
            };

            tracing::info!(
                %path,
                store.size = %bytesize::ByteSize(*store_size).display().si_short(),
                asset.compression = %match compression {
                    Some(c) => c.0.as_str(),
                    None => "none"
                },
                asset.size.source = %bytesize::ByteSize(source_size).display().si_short(),
                asset.size.compressed = compression.is_some().then_some(display(bytesize::ByteSize(size).display().si_short())),
                asset.size.reduction = reduction,
            );
        } else {
            tracing::info!(
                %path,
                compression = %match compression {
                    Some(c) => c.0.as_str(),
                    None => "none",
                }
            );
        }

        drop(store_size);

        Ok(())
    }

    #[instrument(skip_all, err(level = Level::WARN))]
    pub fn get(
        &self,
        path: &str,
        compression: Option<&CompressionAlgorithm>,
    ) -> anyhow::Result<Bytes> {
        tracing::info!(
            path,
            compression = match compression {
                Some(c) => c.as_str(),
                None => "none",
            }
        );

        let mut key = BytesMut::with_capacity(path.len() + 1);
        key.put(path.as_bytes());

        if let Some(compression) = compression {
            key.put_u8(compression.as_u8());
        }

        let txn = ReadTransaction::new(self.env.clone())?;

        let access = txn.access();
        let result = access.get(&self.assets_db, key.as_ref())?;

        Ok(Bytes::copy_from_slice(result))
    }

    #[instrument(skip_all, err)]
    pub fn delete(
        &self,
        path: &str,
        compression: Option<&CompressionAlgorithm>,
    ) -> anyhow::Result<()> {
        tracing::info!(
            path,
            compression = match compression {
                Some(c) => c.as_str(),
                None => "none",
            }
        );

        let mut key = BytesMut::with_capacity(path.len() + 1);
        key.put(path.as_bytes());

        if let Some(compression) = compression {
            key.put_u8(compression.as_u8());
        }

        let txn = WriteTransaction::new(self.env.clone())?;

        {
            let mut access = txn.access();
            let result = access.get::<[u8], [u8]>(&self.assets_db, key.as_ref())?;

            let mut store_size = self.store_size.lock();
            *store_size -= (result.len() + key.len()) as u64;

            drop(store_size);

            access.del_key(&self.assets_db, key.as_ref())?;
        }

        txn.commit()?;

        Ok(())
    }

    #[instrument(skip_all, err)]
    pub fn delete_all(&self, skip_list: &Vec<String>) -> anyhow::Result<()> {
        #[cfg(tracing_unstable)]
        tracing::info!(skip = tracing::field::valuable(&skip_list));
        #[cfg(not(tracing_unstable))]
        tracing::info!(skip = tracing::field::debug(&skip_list));

        let mut skip_set = HashSet::new();

        for skip in skip_list {
            let mut key = BytesMut::with_capacity(skip.len() + 1);
            key.put(skip.as_bytes());

            skip_set.insert(key.clone());

            for compression in [
                CompressionAlgorithm::Brotli,
                CompressionAlgorithm::Deflate,
                CompressionAlgorithm::Gzip,
                CompressionAlgorithm::Zstd { level: 0 },
            ] {
                key.put_u8(compression.as_u8());
                skip_set.insert(key.clone());

                key.truncate(skip.len());
            }
        }

        let mut delete_keys: Vec<(Bytes, usize)> = Vec::new();

        {
            let txn = ReadTransaction::new(self.env.clone())?;

            let access = txn.access();

            let mut cursor = txn.cursor(self.assets_db.clone())?;
            let (key, result) = cursor.first::<[u8], [u8]>(&access)?;

            if !skip_set.contains(key) {
                delete_keys.push((Bytes::copy_from_slice(key), result.len() + key.len()));
            }

            while let Ok((key, result)) = cursor.next::<[u8], [u8]>(&access) {
                if !skip_set.contains(key) {
                    delete_keys.push((Bytes::copy_from_slice(key), result.len() + key.len()));
                }
            }
        }

        let mut store_size = self.store_size.lock();
        let initial_size = *store_size;

        let txn = WriteTransaction::new(self.env.clone())?;

        {
            let mut access = txn.access();

            for (key, size) in delete_keys {
                access.del_key(&self.assets_db, key.as_ref())?;
                *store_size -= size as u64;
            }
        }

        txn.commit()?;

        if self.log_sizes {
            tracing::info!(store.size = %bytesize::ByteSize(*store_size), store.size.reduced = %bytesize::ByteSize(*store_size - initial_size));
        }

        Ok(())
    }
}