extism_runtime/
context.rs

1use std::collections::{BTreeMap, VecDeque};
2
3use crate::*;
4
5static mut TIMER: std::sync::Mutex<Option<Timer>> = std::sync::Mutex::new(None);
6
7/// A `Context` is used to store and manage plugins
8pub struct Context {
9    /// Plugin registry
10    pub plugins: BTreeMap<PluginIndex, Plugin>,
11
12    /// Error message
13    pub error: Option<std::ffi::CString>,
14    next_id: std::sync::atomic::AtomicI32,
15    reclaimed_ids: VecDeque<PluginIndex>,
16
17    // Timeout thread
18    pub(crate) epoch_timer_tx: std::sync::mpsc::SyncSender<TimerAction>,
19}
20
21impl Default for Context {
22    fn default() -> Self {
23        Context::new()
24    }
25}
26
27const START_REUSING_IDS: usize = 25;
28
29impl Context {
30    pub(crate) fn timer() -> std::sync::MutexGuard<'static, Option<Timer>> {
31        match unsafe { TIMER.lock() } {
32            Ok(x) => x,
33            Err(e) => e.into_inner(),
34        }
35    }
36
37    /// Create a new context
38    pub fn new() -> Context {
39        let timer = &mut *Self::timer();
40
41        let tx = match timer {
42            None => Timer::init(timer),
43            Some(t) => t.tx.clone(),
44        };
45
46        Context {
47            plugins: BTreeMap::new(),
48            error: None,
49            next_id: std::sync::atomic::AtomicI32::new(0),
50            reclaimed_ids: VecDeque::new(),
51            epoch_timer_tx: tx,
52        }
53    }
54
55    /// Get the next valid plugin ID
56    pub fn next_id(&mut self) -> Result<PluginIndex, Error> {
57        // Make sure we haven't exhausted all plugin IDs, to reach this it would require the machine
58        // running this code to have a lot of memory - no computer I tested on was able to allocate
59        // the max number of plugins.
60        //
61        // Since `Context::remove` collects IDs that have been removed we will
62        // try to use one of those before returning an error
63        let exhausted = self.next_id.load(std::sync::atomic::Ordering::SeqCst) == PluginIndex::MAX;
64
65        // If there are a significant number of old IDs we can start to re-use them
66        if self.reclaimed_ids.len() >= START_REUSING_IDS || exhausted {
67            if let Some(x) = self.reclaimed_ids.pop_front() {
68                return Ok(x);
69            }
70
71            if exhausted {
72                return Err(anyhow::format_err!(
73                    "All plugin descriptors are in use, unable to allocate a new plugin"
74                ));
75            }
76        }
77
78        Ok(self
79            .next_id
80            .fetch_add(1, std::sync::atomic::Ordering::SeqCst))
81    }
82
83    pub fn insert(&mut self, plugin: Plugin) -> PluginIndex {
84        // Generate a new plugin ID
85        let id: i32 = match self.next_id() {
86            Ok(id) => id,
87            Err(e) => {
88                error!("Error creating Plugin: {:?}", e);
89                self.set_error(e);
90                return -1;
91            }
92        };
93        self.plugins.insert(id, plugin);
94        id
95    }
96
97    pub fn new_plugin<'a>(
98        &mut self,
99        data: impl AsRef<[u8]>,
100        imports: impl IntoIterator<Item = &'a Function>,
101        with_wasi: bool,
102    ) -> PluginIndex {
103        let plugin = match Plugin::new(data, imports, with_wasi) {
104            Ok(x) => x,
105            Err(e) => {
106                error!("Error creating Plugin: {:?}", e);
107                self.set_error(e);
108                return -1;
109            }
110        };
111        self.insert(plugin)
112    }
113
114    /// Set the context error
115    pub fn set_error(&mut self, e: impl std::fmt::Debug) {
116        trace!("Set context error: {:?}", e);
117        self.error = Some(error_string(e));
118    }
119
120    /// Convenience function to set error and return the value passed as the final parameter
121    pub fn error<T>(&mut self, e: impl std::fmt::Debug, x: T) -> T {
122        self.set_error(e);
123        x
124    }
125
126    /// Get a plugin from the context
127    pub fn plugin(&mut self, id: PluginIndex) -> Option<*mut Plugin> {
128        match self.plugins.get_mut(&id) {
129            Some(x) => Some(x),
130            None => None,
131        }
132    }
133
134    pub fn plugin_exists(&mut self, id: PluginIndex) -> bool {
135        self.plugins.contains_key(&id)
136    }
137
138    /// Remove a plugin from the context
139    pub fn remove(&mut self, id: PluginIndex) {
140        if self.plugins.remove(&id).is_some() {
141            // Collect old IDs in case we need to re-use them
142            self.reclaimed_ids.push_back(id);
143        }
144    }
145}
extism_runtime/context.rs

extism_runtime/
context.rs
