termivar_scanner/lua_engine/
registry.rs

1use super::history::BoundedExecutionHistory;
2use super::{
3    LuaExecutionReceipt, LuaExecutionResult, LuaRegistryError, LuaScript, LuaScriptManifest,
4    LuaScriptRegistry, ScriptCategory,
5};
6use crate::lua_config::LuaEngineConfig;
7use std::collections::BTreeMap;
8use std::sync::{Arc, Mutex};
9use tokio::sync::Semaphore;
10
11pub(super) struct RegisteredScript {
12    pub(super) script: LuaScript,
13    pub(super) enabled: bool,
14    pub(super) active_invocations: usize,
15    pub(super) generation: u64,
16}
17
18pub(super) struct RegistryState {
19    pub(super) scripts: BTreeMap<String, RegisteredScript>,
20    names: BTreeMap<String, String>,
21    histories: BTreeMap<String, BoundedExecutionHistory>,
22    history_bytes: usize,
23    pub(super) next_sequence: u64,
24    next_generation: u64,
25    total_source_bytes: usize,
26}
27
28impl RegistryState {
29    fn new() -> Self {
30        Self {
31            scripts: BTreeMap::new(),
32            names: BTreeMap::new(),
33            histories: BTreeMap::new(),
34            history_bytes: 0,
35            next_sequence: 0,
36            next_generation: 0,
37            total_source_bytes: 0,
38        }
39    }
40
41    fn allocate_sequence(&mut self) -> Result<u64, LuaRegistryError> {
42        let sequence = self.next_sequence;
43        self.next_sequence = self
44            .next_sequence
45            .checked_add(1)
46            .ok_or(LuaRegistryError::HistorySequenceExhausted)?;
47        Ok(sequence)
48    }
49
50    fn allocate_generation(&mut self) -> Result<u64, LuaRegistryError> {
51        let generation = self.next_generation;
52        self.next_generation = self
53            .next_generation
54            .checked_add(1)
55            .ok_or(LuaRegistryError::RegistrationGenerationExhausted)?;
56        Ok(generation)
57    }
58
59    fn evict_global_oldest(&mut self) -> bool {
60        let oldest = self
61            .histories
62            .iter()
63            .filter_map(|(script_id, history)| {
64                history
65                    .entries
66                    .front()
67                    .map(|entry| (entry.sequence, script_id.clone()))
68            })
69            .min();
70        let Some((_, script_id)) = oldest else {
71            return false;
72        };
73        let Some(history) = self.histories.get_mut(&script_id) else {
74            return false;
75        };
76        if let Some(entry) = history.pop_front() {
77            self.history_bytes = self.history_bytes.saturating_sub(entry.retained_bytes);
78        }
79        if history.entries.is_empty() {
80            self.histories.remove(&script_id);
81        }
82        true
83    }
84}
85
86pub(super) struct InvocationLease {
87    pub(super) state: Arc<Mutex<RegistryState>>,
88    pub(super) script_id: String,
89    pub(super) generation: u64,
90}
91
92impl Drop for InvocationLease {
93    fn drop(&mut self) {
94        let Ok(mut state) = self.state.lock() else {
95            return;
96        };
97        let Some(entry) = state.scripts.get_mut(&self.script_id) else {
98            return;
99        };
100        if entry.generation == self.generation {
101            entry.active_invocations = entry.active_invocations.saturating_sub(1);
102        }
103    }
104}
105
106impl LuaScriptRegistry {
107    pub fn new() -> Result<Self, LuaRegistryError> {
108        Self::from_config(&LuaEngineConfig::default())
109    }
110
111    pub fn from_config(config: &LuaEngineConfig) -> Result<Self, LuaRegistryError> {
112        config.validate().map_err(LuaRegistryError::InvalidConfig)?;
113        Ok(Self {
114            state: Arc::new(Mutex::new(RegistryState::new())),
115            config: config.clone(),
116            execution_permits: Arc::new(Semaphore::new(config.max_concurrent_executions)),
117        })
118    }
119
120    pub fn register(&self, script: LuaScript) -> Result<(), LuaRegistryError> {
121        let mut state = self
122            .state
123            .lock()
124            .map_err(|_| LuaRegistryError::StateUnavailable)?;
125        let id = script.id();
126        if state.scripts.contains_key(&id) {
127            return Err(LuaRegistryError::DuplicateId);
128        }
129        if state.names.contains_key(&script.name) {
130            return Err(LuaRegistryError::DuplicateName);
131        }
132        if state.scripts.len() >= self.config.max_scripts {
133            return Err(LuaRegistryError::ScriptCapacity);
134        }
135        if script.source.len() > self.config.max_source_bytes {
136            return Err(LuaRegistryError::SourceLimit);
137        }
138        let total_source_bytes = state
139            .total_source_bytes
140            .checked_add(script.source.len())
141            .ok_or(LuaRegistryError::TotalSourceCapacity)?;
142        if total_source_bytes > self.config.max_total_source_bytes {
143            return Err(LuaRegistryError::TotalSourceCapacity);
144        }
145        let generation = state.allocate_generation()?;
146        state.names.insert(script.name.clone(), id.clone());
147        state.scripts.insert(
148            id,
149            RegisteredScript {
150                enabled: script.enabled,
151                script,
152                active_invocations: 0,
153                generation,
154            },
155        );
156        state.total_source_bytes = total_source_bytes;
157        Ok(())
158    }
159
160    pub fn get(&self, script_id: &str) -> Result<Option<LuaScriptManifest>, LuaRegistryError> {
161        let state = self
162            .state
163            .lock()
164            .map_err(|_| LuaRegistryError::StateUnavailable)?;
165        Ok(state
166            .scripts
167            .get(script_id)
168            .map(|entry| entry.script.manifest_with_enabled(entry.enabled)))
169    }
170
171    pub fn list_all(&self) -> Result<Vec<LuaScriptManifest>, LuaRegistryError> {
172        let state = self
173            .state
174            .lock()
175            .map_err(|_| LuaRegistryError::StateUnavailable)?;
176        let mut manifests: Vec<_> = state
177            .scripts
178            .values()
179            .map(|entry| entry.script.manifest_with_enabled(entry.enabled))
180            .collect();
181        manifests.sort_by(|left, right| {
182            (&left.name, &left.version, &left.id).cmp(&(&right.name, &right.version, &right.id))
183        });
184        Ok(manifests)
185    }
186
187    pub fn list_enabled(&self) -> Result<Vec<LuaScriptManifest>, LuaRegistryError> {
188        Ok(self
189            .list_all()?
190            .into_iter()
191            .filter(LuaScriptManifest::enabled)
192            .collect())
193    }
194
195    pub fn list_by_category(
196        &self,
197        category: ScriptCategory,
198    ) -> Result<Vec<LuaScriptManifest>, LuaRegistryError> {
199        Ok(self
200            .list_all()?
201            .into_iter()
202            .filter(|manifest| manifest.categories.contains(&category))
203            .collect())
204    }
205
206    pub fn get_history(
207        &self,
208        script_id: &str,
209    ) -> Result<Vec<LuaExecutionReceipt>, LuaRegistryError> {
210        let state = self
211            .state
212            .lock()
213            .map_err(|_| LuaRegistryError::StateUnavailable)?;
214        Ok(state
215            .histories
216            .get(script_id)
217            .map(|history| {
218                history
219                    .entries
220                    .iter()
221                    .map(|entry| entry.receipt.clone())
222                    .collect()
223            })
224            .unwrap_or_default())
225    }
226
227    pub fn get_recent_history(
228        &self,
229        script_id: &str,
230        count: usize,
231    ) -> Result<Vec<LuaExecutionReceipt>, LuaRegistryError> {
232        let state = self
233            .state
234            .lock()
235            .map_err(|_| LuaRegistryError::StateUnavailable)?;
236        Ok(state
237            .histories
238            .get(script_id)
239            .map(|history| {
240                history
241                    .entries
242                    .iter()
243                    .rev()
244                    .take(count)
245                    .map(|entry| entry.receipt.clone())
246                    .collect()
247            })
248            .unwrap_or_default())
249    }
250
251    pub fn count(&self) -> Result<usize, LuaRegistryError> {
252        self.state
253            .lock()
254            .map(|state| state.scripts.len())
255            .map_err(|_| LuaRegistryError::StateUnavailable)
256    }
257
258    pub fn enabled_count(&self) -> Result<usize, LuaRegistryError> {
259        self.state
260            .lock()
261            .map(|state| state.scripts.values().filter(|entry| entry.enabled).count())
262            .map_err(|_| LuaRegistryError::StateUnavailable)
263    }
264
265    pub fn set_enabled(&self, script_id: &str, enabled: bool) -> Result<(), LuaRegistryError> {
266        let mut state = self
267            .state
268            .lock()
269            .map_err(|_| LuaRegistryError::StateUnavailable)?;
270        let entry = state
271            .scripts
272            .get_mut(script_id)
273            .ok_or(LuaRegistryError::ScriptNotFound)?;
274        entry.enabled = enabled;
275        Ok(())
276    }
277
278    pub fn unregister(&self, script_id: &str) -> Result<(), LuaRegistryError> {
279        let mut state = self
280            .state
281            .lock()
282            .map_err(|_| LuaRegistryError::StateUnavailable)?;
283        let entry = state
284            .scripts
285            .get(script_id)
286            .ok_or(LuaRegistryError::ScriptNotFound)?;
287        if entry.active_invocations != 0 {
288            return Err(LuaRegistryError::ScriptInUse);
289        }
290        let total_source_bytes = state
291            .total_source_bytes
292            .checked_sub(entry.script.source.len())
293            .ok_or(LuaRegistryError::StateUnavailable)?;
294        let removed_history_bytes = state
295            .histories
296            .get(script_id)
297            .map_or(0, |history| history.retained_bytes);
298        let history_bytes = state
299            .history_bytes
300            .checked_sub(removed_history_bytes)
301            .ok_or(LuaRegistryError::StateUnavailable)?;
302        let entry = state
303            .scripts
304            .remove(script_id)
305            .ok_or(LuaRegistryError::ScriptNotFound)?;
306        state.names.remove(&entry.script.name);
307        state.total_source_bytes = total_source_bytes;
308        state.histories.remove(script_id);
309        state.history_bytes = history_bytes;
310        Ok(())
311    }
312
313    pub(super) fn record_result(
314        &self,
315        generation: u64,
316        result: &LuaExecutionResult,
317    ) -> Result<(), LuaRegistryError> {
318        let mut state = self
319            .state
320            .lock()
321            .map_err(|_| LuaRegistryError::StateUnavailable)?;
322        if state
323            .scripts
324            .get(result.script_id())
325            .is_none_or(|entry| entry.generation != generation)
326        {
327            return Ok(());
328        }
329        let receipt = LuaExecutionReceipt::from_result(result);
330        let receipt_bytes = receipt.retained_bytes();
331        if receipt_bytes > self.config.max_history_bytes_per_script
332            || receipt_bytes > self.config.max_history_bytes_total
333        {
334            return Ok(());
335        }
336        let sequence = state.allocate_sequence()?;
337        while state.history_bytes.saturating_add(receipt_bytes)
338            > self.config.max_history_bytes_total
339        {
340            if !state.evict_global_oldest() {
341                return Ok(());
342            }
343        }
344        let history = state
345            .histories
346            .entry(result.script_id().to_owned())
347            .or_insert_with(BoundedExecutionHistory::new);
348        let before = history.retained_bytes;
349        let inserted = history.push(
350            sequence,
351            receipt,
352            self.config.history_size,
353            self.config.max_history_bytes_per_script,
354        );
355        let after = history.retained_bytes;
356        if inserted {
357            state.history_bytes = state
358                .history_bytes
359                .saturating_sub(before)
360                .saturating_add(after);
361        }
362        Ok(())
363    }
364}
termivar_scanner/lua_engine/registry.rs

termivar_scanner/lua_engine/
registry.rs
