1
use std::sync::Arc;
2
use std::sync::atomic::{AtomicU64, Ordering};
3

            
4
use uuid::Uuid;
5
use wasmtime::Engine;
6

            
7
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8
pub struct ScriptLimits {
9
    pub fuel: u64,
10
    pub max_memory_pages: u64,
11
}
12

            
13
impl Default for ScriptLimits {
14
1698
    fn default() -> Self {
15
1698
        Self {
16
1698
            fuel: 1_000_000,
17
1698
            max_memory_pages: 64,
18
1698
        }
19
1698
    }
20
}
21

            
22
#[derive(Debug, Clone)]
23
pub struct ScriptCtx {
24
    pub user_id: Uuid,
25
    pub limits: ScriptLimits,
26
}
27

            
28
impl ScriptCtx {
29
    #[must_use]
30
1643
    pub fn new(user_id: Uuid) -> Self {
31
1643
        Self {
32
1643
            user_id,
33
1643
            limits: ScriptLimits::default(),
34
1643
        }
35
1643
    }
36

            
37
    #[must_use]
38
73
    pub fn with_limits(mut self, limits: ScriptLimits) -> Self {
39
73
        self.limits = limits;
40
73
        self
41
73
    }
42
}
43

            
44
/// Cooperative interrupt signal shared between the SLYNK reader task (which
45
/// `interrupt()`s on `(:emacs-interrupt)`) and the eval task. Backed by a
46
/// monotonic generation counter rather than a boolean latch: an interrupt
47
/// increments the count, and a request asks "did the count move since I
48
/// started?" via [`InterruptHandle::generation`] + [`InterruptHandle::is_interrupted_since`].
49
///
50
/// The generation design is deliberate (a prior boolean latch was racy): there
51
/// is no shared flag to leave set, so an interrupt can never "stick" and poison
52
/// an unrelated later request. Each top-level operation snapshots the
53
/// generation at entry and compares against it; the next operation re-snapshots,
54
/// automatically discarding any interrupt that targeted the finished one — no
55
/// per-exit cleanup needed, on any terminal path.
56
#[derive(Debug, Clone, Default)]
57
pub struct InterruptHandle {
58
    generation: Arc<AtomicU64>,
59
}
60

            
61
impl InterruptHandle {
62
1541
    pub(crate) fn new() -> Self {
63
1541
        Self::default()
64
1541
    }
65

            
66
    /// Signals an interrupt by advancing the generation. Idempotent in effect:
67
    /// any request whose baseline predates this call observes the interrupt.
68
121
    pub fn interrupt(&self) {
69
121
        self.generation.fetch_add(1, Ordering::SeqCst);
70
121
    }
71

            
72
    /// The current interrupt generation. The eval task keeps an "ack"
73
    /// watermark of the highest generation already attributed to a finished
74
    /// request; a new interrupt is one whose generation exceeds that watermark.
75
5391
    pub(crate) fn generation(&self) -> u64 {
76
5391
        self.generation.load(Ordering::SeqCst)
77
5391
    }
78
}
79

            
80
/// Cooperative cancel handle for an in-flight eval. Wraps a clone of
81
/// the Session's [`wasmtime::Engine`]; the engine's epoch counter is
82
/// shared internally so a `.bump()` from any task takes effect on the
83
/// concurrent `func.call_async` running against the same engine.
84
///
85
/// Each Session sets `store.set_epoch_deadline(1)` before invoking
86
/// `nomi-eval`, so one bump trips the deadline and the eval traps
87
/// with `wasmtime::Trap::Interrupt` — surfaced upstream as
88
/// `EngineError::EpochInterrupt` and on the wire as
89
/// `(:error (:code interrupted ...))`.
90
#[derive(Clone, Debug)]
91
pub struct EpochBumper {
92
    engine: Engine,
93
}
94

            
95
impl EpochBumper {
96
200
    pub(crate) fn new(engine: Engine) -> Self {
97
200
        Self { engine }
98
200
    }
99

            
100
    /// Increments the engine epoch. Returns immediately (no await).
101
    /// Safe to call before, during, or after an eval; the next epoch
102
    /// check inside any in-flight `call_async` traps.
103
92
    pub fn bump(&self) {
104
92
        self.engine.increment_epoch();
105
92
    }
106
}
107

            
108
#[cfg(test)]
109
mod tests {
110
    use super::*;
111

            
112
    #[test]
113
1
    fn default_limits_are_non_zero() {
114
1
        let limits = ScriptLimits::default();
115
1
        assert!(limits.fuel > 0);
116
1
        assert!(limits.max_memory_pages > 0);
117
1
    }
118

            
119
    #[test]
120
1
    fn ctx_carries_user_id() {
121
1
        let id = Uuid::new_v4();
122
1
        let ctx = ScriptCtx::new(id);
123
1
        assert_eq!(ctx.user_id, id);
124
1
    }
125

            
126
    #[test]
127
1
    fn with_limits_overrides_default() {
128
1
        let id = Uuid::new_v4();
129
1
        let limits = ScriptLimits {
130
1
            fuel: 42,
131
1
            max_memory_pages: 7,
132
1
        };
133
1
        let ctx = ScriptCtx::new(id).with_limits(limits);
134
1
        assert_eq!(ctx.limits, limits);
135
1
    }
136

            
137
    #[test]
138
1
    fn interrupt_advances_generation_monotonically() {
139
1
        let h = InterruptHandle::new();
140
1
        let base = h.generation();
141
1
        h.interrupt();
142
1
        assert_eq!(h.generation(), base + 1);
143
1
        h.interrupt();
144
1
        assert_eq!(h.generation(), base + 2);
145
1
    }
146

            
147
    #[test]
148
1
    fn interrupt_handle_clones_share_generation() {
149
1
        let a = InterruptHandle::new();
150
1
        let base = a.generation();
151
1
        let b = a.clone();
152
1
        b.interrupt();
153
1
        assert!(a.generation() > base);
154
1
    }
155
}