ordinary-storage 0.11.1

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

use anyhow::bail;
use bytes::{BufMut, Bytes, BytesMut};
use ordinary_config::ArtifactLimits;
use parking_lot::Mutex;
use saferlmdb::{
    self as lmdb, Database, DatabaseOptions, Environment, ReadTransaction, WriteTransaction, put,
};
use std::sync::Arc;
use tracing::Level;
use tracing::instrument;

#[derive(Debug)]
pub enum ArtifactKind {
    FunctionWasiP2,
    ExtensionWasiP2,
}

impl ArtifactKind {
    fn key(&self, name: &[u8]) -> Bytes {
        let mut key = BytesMut::new();

        match self {
            ArtifactKind::FunctionWasiP2 => key.put_u8(3),
            ArtifactKind::ExtensionWasiP2 => key.put_u8(4),
        }

        key.put(name);
        key.into()
    }
}

pub struct ArtifactStore {
    pub limits: ArtifactLimits,
    env: Arc<Environment>,

    /// stores artifacts to be executed by plugin and template tasks
    artifact_db: Arc<Database<'static>>,

    log_sizes: bool,

    store_size: Arc<Mutex<u64>>,
}

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

        let mut store_size = 0;

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

        let mut artifact_cursor = txn.cursor(artifact_db.clone())?;

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

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

        Ok(Self {
            limits,
            env: env.clone(),
            artifact_db,
            log_sizes,
            store_size: Arc::new(Mutex::new(store_size)),
        })
    }

    #[instrument(skip_all, err)]
    pub fn put(&self, name: &str, kind: ArtifactKind, content: &[u8]) -> anyhow::Result<()> {
        let key = kind.key(name.as_bytes());

        let compressed = zstd::stream::encode_all(std::io::Cursor::new(content), 17)?;
        let size = (key.len() + compressed.len()) as u64;

        if size > self.limits.max_artifact_size {
            bail!("exceeds max artifact size");
        }

        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.artifact_db, key.as_ref()) {
                overwrite_size = (key.len() + result.len()) as u64;
            }

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

        txn.commit()?;

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

        if self.log_sizes {
            tracing::info!(
                size.source = %bytesize::ByteSize((key.len() + content.len()) as u64).display().si_short(),
                size.compressed = %bytesize::ByteSize((key.len() + compressed.len()) as u64).display().si_short(),
                size.stored = %bytesize::ByteSize(*store_size).display().si_short(),
                "size"
            );
        }

        drop(store_size);

        Ok(())
    }

    #[instrument(skip_all, err(level = Level::WARN))]
    pub fn get(&self, name: &str, kind: ArtifactKind) -> anyhow::Result<Bytes> {
        let key = kind.key(name.as_bytes());
        let txn = ReadTransaction::new(self.env.clone())?;

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

        let decompressed = zstd::stream::decode_all(result)?;

        Ok(Bytes::copy_from_slice(&decompressed))
    }

    #[instrument(skip_all, err)]
    pub fn delete(&self, name: &str, kind: ArtifactKind) -> anyhow::Result<()> {
        let key = kind.key(name.as_bytes());
        let txn = WriteTransaction::new(self.env.clone())?;

        {
            let mut access = txn.access();
            let result = access.get::<[u8], [u8]>(&self.artifact_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.artifact_db, key.as_ref())?;
        }

        txn.commit()?;

        Ok(())
    }
}