Skip to main content

byteflow/jit/
runtime.rs

1//! Shared JIT state for the M:N runtime (cache + hot counters, lock-free execution path).
2
3use std::sync::{Arc, Mutex, RwLock};
4
5use crate::Chunk;
6
7use super::trace::{HotCounter, TraceCache, TraceKey, HOT_THRESHOLD};
8use super::CompiledTrace;
9
10/// Runtime-wide JIT state: compiled traces are shared; each worker thread owns
11/// its own Cranelift module for compilation ([`super::module_local`]).
12pub struct JitRuntime {
13    chunk: RwLock<Arc<Chunk>>,
14    pub hot_threshold: u32,
15    cache: RwLock<TraceCache>,
16    hot: Mutex<HotCounter>,
17}
18
19impl JitRuntime {
20    pub fn new(chunk: Arc<Chunk>, hot_threshold: u32) -> Self {
21        Self {
22            chunk: RwLock::new(chunk),
23            hot_threshold,
24            cache: RwLock::new(TraceCache::new()),
25            hot: Mutex::new(HotCounter::default()),
26        }
27    }
28
29    pub fn with_default_threshold(chunk: Arc<Chunk>) -> Self {
30        Self::new(chunk, HOT_THRESHOLD)
31    }
32
33    /// Current bytecode image this cache was built for.
34    pub fn chunk(&self) -> Option<Arc<Chunk>> {
35        self.chunk.read().ok().map(|g| Arc::clone(&g))
36    }
37
38    /// True when `vm_chunk` is the same image as the JIT cache.
39    pub fn matches_chunk(&self, vm_chunk: &Arc<Chunk>) -> bool {
40        match self.chunk.read() {
41            Ok(guard) => Arc::ptr_eq(&guard, vm_chunk),
42            Err(_) => false,
43        }
44    }
45
46    /// Replace the bytecode image and discard stale compiled traces.
47    pub fn reload(&self, chunk: Arc<Chunk>) {
48        if let Ok(mut current) = self.chunk.write() {
49            *current = chunk;
50        }
51        if let Ok(mut cache) = self.cache.write() {
52            cache.clear();
53        }
54        if let Ok(mut hot) = self.hot.lock() {
55            hot.clear();
56        }
57    }
58
59    pub fn get_trace(&self, key: &TraceKey) -> Option<JitEntryCopy> {
60        let cache = self.cache.read().ok()?;
61        cache.get(key).map(|t| JitEntryCopy {
62            entry: t.entry,
63        })
64    }
65
66    pub fn insert_trace(&self, key: TraceKey, trace: CompiledTrace) {
67        if let Ok(mut cache) = self.cache.write() {
68            cache.insert(key, trace);
69        }
70    }
71
72    pub fn record_hot_hit(&self, key: TraceKey) -> bool {
73        let Ok(mut hot) = self.hot.lock() else {
74            return false;
75        };
76        hot.hit(key, self.hot_threshold)
77    }
78}
79
80/// [`JitEntry`] is a function pointer and safe to copy out of the cache briefly.
81#[derive(Clone, Copy)]
82pub struct JitEntryCopy {
83    pub entry: super::frame::JitEntry,
84}