Skip to main content

ordinary_storage/stores/
artifact.rs

1// Copyright (C) 2026 Ordinary Labs, LLC.
2//
3// SPDX-License-Identifier: AGPL-3.0-only
4
5use anyhow::bail;
6use bytes::Bytes;
7use ordinary_config::ArtifactLimits;
8use parking_lot::Mutex;
9use saferlmdb::{
10    self as lmdb, Database, DatabaseOptions, Environment, ReadTransaction, WriteTransaction, put,
11};
12use std::sync::Arc;
13use tracing::Level;
14use tracing::instrument;
15
16#[derive(Debug)]
17pub enum ArtifactKind {
18    Template,
19    Action,
20}
21
22pub struct ArtifactStore {
23    pub limits: ArtifactLimits,
24    env: Arc<Environment>,
25
26    /// stores artifacts to be executed by plugin and template tasks
27    artifact_db: Arc<Database<'static>>,
28
29    log_size: bool,
30
31    store_size: Arc<Mutex<u64>>,
32}
33
34impl ArtifactStore {
35    pub fn new(
36        limits: ArtifactLimits,
37        env: &Arc<Environment>,
38        log_size: bool,
39    ) -> anyhow::Result<Self> {
40        let artifact_db = Arc::new(Database::open(
41            env.clone(),
42            Some("artifact"),
43            &DatabaseOptions::new(lmdb::db::Flags::CREATE),
44        )?);
45
46        let mut store_size = 0;
47
48        let txn = ReadTransaction::new(env.clone())?;
49        let access = txn.access();
50
51        let mut artifact_cursor = txn.cursor(artifact_db.clone())?;
52
53        if let Ok((key, val)) = artifact_cursor.first::<[u8], [u8]>(&access) {
54            store_size += key.len() as u64;
55            store_size += val.len() as u64;
56
57            while let Ok((key, val)) = artifact_cursor.next::<[u8], [u8]>(&access) {
58                store_size += key.len() as u64;
59                store_size += val.len() as u64;
60            }
61        }
62
63        Ok(Self {
64            limits,
65            env: env.clone(),
66            artifact_db,
67            log_size,
68            store_size: Arc::new(Mutex::new(store_size)),
69        })
70    }
71
72    #[instrument(skip_all, err)]
73    pub fn put(&self, idx: u8, kind: ArtifactKind, content: &[u8]) -> anyhow::Result<()> {
74        let key = match kind {
75            ArtifactKind::Action => [0, idx],
76            ArtifactKind::Template => [1, idx],
77        };
78
79        let compressed = zstd::stream::encode_all(std::io::Cursor::new(content), 17)?;
80        let size = (key.len() + compressed.len()) as u64;
81
82        if size > self.limits.max_artifact_size {
83            bail!("exceeds max artifact size");
84        }
85
86        let mut store_size = self.store_size.lock();
87
88        if *store_size + size > self.limits.max_store_size {
89            bail!("exceeds store size limit");
90        }
91
92        let mut overwrite_size = 0;
93
94        let txn = WriteTransaction::new(self.env.clone())?;
95
96        {
97            let mut access = txn.access();
98            if let Ok(result) = access.get::<[u8], [u8]>(&self.artifact_db, key.as_ref()) {
99                overwrite_size = (key.len() + result.len()) as u64;
100            }
101
102            access.put(&self.artifact_db, &key, &compressed, &put::Flags::empty())?;
103        }
104
105        txn.commit()?;
106
107        *store_size -= overwrite_size;
108        *store_size += size;
109
110        if self.log_size {
111            tracing::info!(
112                size.source = %bytesize::ByteSize((key.len() + content.len()) as u64).display().si_short(),
113                size.compressed = %bytesize::ByteSize((key.len() + compressed.len()) as u64).display().si_short(),
114                size.stored = %bytesize::ByteSize(*store_size).display().si_short(),
115                "size"
116            );
117        }
118
119        drop(store_size);
120
121        Ok(())
122    }
123
124    #[instrument(skip_all, err(level = Level::WARN))]
125    pub fn get(&self, idx: u8, kind: ArtifactKind) -> anyhow::Result<Bytes> {
126        let key = match kind {
127            ArtifactKind::Action => [0, idx],
128            ArtifactKind::Template => [1, idx],
129        };
130
131        let txn = ReadTransaction::new(self.env.clone())?;
132
133        let access = txn.access();
134        let result = access.get::<[u8], [u8]>(&self.artifact_db, &key)?;
135
136        let decompressed = zstd::stream::decode_all(result)?;
137
138        Ok(Bytes::copy_from_slice(&decompressed))
139    }
140
141    #[instrument(skip_all, err)]
142    pub fn delete(&self, idx: u8, kind: ArtifactKind) -> anyhow::Result<()> {
143        let key = match kind {
144            ArtifactKind::Action => [0, idx],
145            ArtifactKind::Template => [1, idx],
146        };
147
148        let txn = WriteTransaction::new(self.env.clone())?;
149
150        {
151            let mut access = txn.access();
152            let result = access.get::<[u8], [u8]>(&self.artifact_db, key.as_ref())?;
153
154            let mut store_size = self.store_size.lock();
155            *store_size -= (result.len() + key.len()) as u64;
156
157            drop(store_size);
158
159            access.del_key(&self.artifact_db, &key)?;
160        }
161
162        txn.commit()?;
163
164        Ok(())
165    }
166}