memlink_runtime/
reload.rs

1//! Hot-reload support for module reloading.
2
3use std::sync::Arc;
4use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
5use std::time::Duration;
6
7use crate::error::{Error, Result};
8use crate::instance::ModuleInstance;
9
10pub const DEFAULT_DRAIN_TIMEOUT: Duration = Duration::from_secs(30);
11
12#[derive(Debug)]
13pub struct ReloadState {
14    pub old_handle: u64,
15    pub new_handle: u64,
16    pub draining: bool,
17    in_flight: Arc<AtomicUsize>,
18    drain_complete: Arc<AtomicBool>,
19}
20
21impl ReloadState {
22    pub fn new(old_handle: u64, new_handle: u64, in_flight_count: usize) -> Self {
23        let in_flight = Arc::new(AtomicUsize::new(in_flight_count));
24        let drain_complete = Arc::new(AtomicBool::new(false));
25
26        if in_flight_count == 0 {
27            drain_complete.store(true, Ordering::Relaxed);
28        }
29
30        ReloadState {
31            old_handle,
32            new_handle,
33            draining: true,
34            in_flight,
35            drain_complete,
36        }
37    }
38
39    pub fn in_flight_count(&self) -> usize {
40        self.in_flight.load(Ordering::Relaxed)
41    }
42
43    pub fn is_drain_complete(&self) -> bool {
44        self.drain_complete.load(Ordering::Relaxed)
45    }
46
47    pub fn mark_call_started(&self) {
48        self.in_flight.fetch_add(1, Ordering::Relaxed);
49        self.drain_complete.store(false, Ordering::Relaxed);
50    }
51
52    pub fn mark_call_completed(&self) {
53        if self.in_flight.fetch_sub(1, Ordering::Relaxed) == 1 {
54            self.drain_complete.store(true, Ordering::Relaxed);
55        }
56    }
57
58    pub fn wait_for_drain(&self, timeout: Duration) -> Result<()> {
59        let start = std::time::Instant::now();
60
61        while start.elapsed() < timeout {
62            if self.is_drain_complete() {
63                return Ok(());
64            }
65            std::thread::sleep(Duration::from_millis(10));
66        }
67
68        Err(Error::ReloadTimeout(self.in_flight_count()))
69    }
70}
71
72#[derive(Debug, Clone)]
73pub struct ModuleState {
74    pub data: Vec<u8>,
75    pub abi_version: u32,
76}
77
78impl ModuleState {
79    pub fn new(data: Vec<u8>, abi_version: u32) -> Self {
80        ModuleState { data, abi_version }
81    }
82}
83
84pub trait StatefulModule {
85    fn serialize_state(&self) -> Result<Vec<u8>>;
86    fn restore_state(&self, _state: &[u8]) -> Result<()>;
87}
88
89impl StatefulModule for ModuleInstance {
90    fn serialize_state(&self) -> Result<Vec<u8>> {
91        Err(Error::UnsupportedReference)
92    }
93
94    fn restore_state(&self, _state: &[u8]) -> Result<()> {
95        Err(Error::UnsupportedReference)
96    }
97}
98
99#[derive(Debug, Clone)]
100pub struct ReloadConfig {
101    pub drain_timeout: Duration,
102    pub preserve_state: bool,
103    pub force_unload_on_timeout: bool,
104}
105
106impl Default for ReloadConfig {
107    fn default() -> Self {
108        ReloadConfig {
109            drain_timeout: DEFAULT_DRAIN_TIMEOUT,
110            preserve_state: false,
111            force_unload_on_timeout: true,
112        }
113    }
114}
115
116impl ReloadConfig {
117    pub fn with_drain_timeout(mut self, timeout: Duration) -> Self {
118        self.drain_timeout = timeout;
119        self
120    }
121
122    pub fn with_state_preservation(mut self) -> Self {
123        self.preserve_state = true;
124        self
125    }
126
127    pub fn with_force_unload(mut self, force: bool) -> Self {
128        self.force_unload_on_timeout = force;
129        self
130    }
131}
memlink_runtime/reload.rs

memlink_runtime/
reload.rs
