termivar_scanner/lua_engine/
execution.rs

1use super::history::ExecutionProvenance;
2use super::registry::InvocationLease;
3use super::vm::execute_snapshot;
4use super::{
5    LuaCancellationToken, LuaContext, LuaExecutionError, LuaExecutionResult, LuaRegistryError,
6    LuaScript, LuaScriptRegistry,
7};
8use std::future::{poll_fn, Future};
9use std::panic::{catch_unwind, AssertUnwindSafe};
10use std::pin::Pin;
11use std::sync::Arc;
12use std::task::Poll;
13use std::time::Instant;
14
15struct RegisteredSnapshot {
16    script: LuaScript,
17    enabled: bool,
18    generation: u64,
19}
20
21impl LuaScriptRegistry {
22    pub async fn execute(
23        &self,
24        script_id: &str,
25        context: LuaContext,
26    ) -> Result<LuaExecutionResult, LuaRegistryError> {
27        self.execute_with_cancellation(script_id, context, LuaCancellationToken::new())
28            .await
29    }
30
31    pub async fn execute_with_cancellation(
32        &self,
33        script_id: &str,
34        context: LuaContext,
35        cancellation: LuaCancellationToken,
36    ) -> Result<LuaExecutionResult, LuaRegistryError> {
37        let snapshot = {
38            let state = self
39                .state
40                .lock()
41                .map_err(|_| LuaRegistryError::StateUnavailable)?;
42            let entry = state
43                .scripts
44                .get(script_id)
45                .ok_or(LuaRegistryError::ScriptNotFound)?;
46            RegisteredSnapshot {
47                script: entry.script.clone(),
48                enabled: entry.enabled,
49                generation: entry.generation,
50            }
51        };
52        let started = Instant::now();
53        if !snapshot.enabled {
54            let result = LuaExecutionResult::failed(
55                ExecutionProvenance::from(&snapshot.script),
56                LuaExecutionError::ScriptDisabled,
57                started,
58            );
59            self.record_result(snapshot.generation, &result)?;
60            return Ok(result);
61        }
62        if let Err(error) = context.validate(&self.config) {
63            let result = LuaExecutionResult::failed(
64                ExecutionProvenance::from(&snapshot.script),
65                error,
66                started,
67            );
68            self.record_result(snapshot.generation, &result)?;
69            return Ok(result);
70        }
71        if cancellation.is_cancelled() {
72            let result = LuaExecutionResult::failed(
73                ExecutionProvenance::from(&snapshot.script),
74                LuaExecutionError::Cancelled,
75                started,
76            );
77            self.record_result(snapshot.generation, &result)?;
78            return Ok(result);
79        }
80        let runtime = match tokio::runtime::Handle::try_current() {
81            Ok(runtime) => runtime,
82            Err(_) => {
83                let result = LuaExecutionResult::failed(
84                    ExecutionProvenance::from(&snapshot.script),
85                    LuaExecutionError::HostFailure,
86                    started,
87                );
88                self.record_result(snapshot.generation, &result)?;
89                return Ok(result);
90            },
91        };
92        let permit = match Arc::clone(&self.execution_permits).try_acquire_owned() {
93            Ok(permit) => permit,
94            Err(_) => {
95                let result = LuaExecutionResult::failed(
96                    ExecutionProvenance::from(&snapshot.script),
97                    LuaExecutionError::ConcurrencyLimit,
98                    started,
99                );
100                self.record_result(snapshot.generation, &result)?;
101                return Ok(result);
102            },
103        };
104        let (script, lease, generation) = {
105            let mut state = self
106                .state
107                .lock()
108                .map_err(|_| LuaRegistryError::StateUnavailable)?;
109            let entry = state
110                .scripts
111                .get_mut(script_id)
112                .ok_or(LuaRegistryError::ScriptNotFound)?;
113            if !entry.enabled {
114                let provenance = ExecutionProvenance::from(&entry.script);
115                let generation = entry.generation;
116                drop(state);
117                drop(permit);
118                let result = LuaExecutionResult::failed(
119                    provenance,
120                    LuaExecutionError::ScriptDisabled,
121                    started,
122                );
123                self.record_result(generation, &result)?;
124                return Ok(result);
125            }
126            entry.active_invocations = entry
127                .active_invocations
128                .checked_add(1)
129                .ok_or(LuaRegistryError::InvocationLimit)?;
130            let script = entry.script.clone();
131            let generation = entry.generation;
132            let lease = InvocationLease {
133                state: Arc::clone(&self.state),
134                script_id: script.id(),
135                generation,
136            };
137            (script, lease, generation)
138        };
139        let config = self.config.clone();
140        let fallback_provenance = ExecutionProvenance::from(&script);
141        let worker = match catch_unwind(AssertUnwindSafe(|| {
142            runtime.spawn_blocking(move || {
143                let _permit = permit;
144                let result = execute_snapshot(script, context, config, cancellation, started);
145                (result, lease)
146            })
147        })) {
148            Ok(worker) => worker,
149            Err(_) => {
150                let result = LuaExecutionResult::failed(
151                    fallback_provenance,
152                    LuaExecutionError::HostFailure,
153                    started,
154                );
155                self.record_result(generation, &result)?;
156                return Ok(result);
157            },
158        };
159        let result = match await_worker(worker).await {
160            Ok((result, lease)) => {
161                self.record_result(generation, &result)?;
162                drop(lease);
163                result
164            },
165            Err(_) => {
166                let result = LuaExecutionResult::failed(
167                    fallback_provenance,
168                    LuaExecutionError::HostFailure,
169                    started,
170                );
171                self.record_result(generation, &result)?;
172                result
173            },
174        };
175        Ok(result)
176    }
177}
178async fn await_worker<T>(mut worker: tokio::task::JoinHandle<T>) -> Result<T, ()> {
179    poll_fn(|context| {
180        match catch_unwind(AssertUnwindSafe(|| Pin::new(&mut worker).poll(context))) {
181            Ok(Poll::Ready(result)) => Poll::Ready(result.map_err(|_| ())),
182            Ok(Poll::Pending) => Poll::Pending,
183            Err(_) => {
184                worker.abort();
185                Poll::Ready(Err(()))
186            },
187        }
188    })
189    .await
190}
termivar_scanner/lua_engine/execution.rs

termivar_scanner/lua_engine/
execution.rs
