ordinary_storage/stores/
artifact.rs1use 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 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}