Skip to main content

ordinary_storage/stores/
artifact.rs

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