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