posemesh_node_registration/
state.rs

1use anyhow::{anyhow, Result};
2use chrono::{DateTime, Utc};
3use serde::{Deserialize, Serialize};
4use std::io;
5use std::sync::{Mutex, MutexGuard, OnceLock};
6use std::time::{Duration, Instant};
7
8pub const STATUS_REGISTERING: &str = "registering";
9pub const STATUS_REGISTERED: &str = "registered";
10pub const STATUS_DISCONNECTED: &str = "disconnected";
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct RegistrationState {
14    pub status: String,
15    #[serde(default, with = "chrono::serde::ts_seconds_option")]
16    pub last_healthcheck: Option<DateTime<Utc>>,
17}
18
19impl Default for RegistrationState {
20    fn default() -> Self {
21        Self {
22            status: STATUS_DISCONNECTED.to_string(),
23            last_healthcheck: None,
24        }
25    }
26}
27
28#[derive(Default)]
29struct InMemoryState {
30    registration: Option<RegistrationState>,
31    node_secret: Option<String>,
32}
33
34static STATE_STORE: OnceLock<Mutex<InMemoryState>> = OnceLock::new();
35
36fn state_store() -> &'static Mutex<InMemoryState> {
37    STATE_STORE.get_or_init(|| Mutex::new(InMemoryState::default()))
38}
39
40fn lock_state_store() -> Result<MutexGuard<'static, InMemoryState>> {
41    state_store()
42        .lock()
43        .map_err(|_| anyhow!("registration state store poisoned"))
44}
45
46/// Store node secret bytes in memory.
47pub fn write_node_secret(secret: &str) -> Result<()> {
48    let mut store = lock_state_store()?;
49    store.node_secret = Some(secret.to_owned());
50    Ok(())
51}
52
53/// Read secret contents. Returns Ok(None) if missing.
54pub fn read_node_secret() -> Result<Option<String>> {
55    let store = lock_state_store()?;
56    Ok(store.node_secret.clone())
57}
58
59/// Clear any stored secret. Intended for tests.
60pub fn clear_node_secret() -> Result<()> {
61    let mut store = lock_state_store()?;
62    store.node_secret = None;
63    Ok(())
64}
65
66pub fn read_state() -> Result<RegistrationState> {
67    let store = lock_state_store()?;
68    Ok(store.registration.clone().unwrap_or_default())
69}
70
71pub fn write_state(st: &RegistrationState) -> Result<()> {
72    let mut store = lock_state_store()?;
73    store.registration = Some(st.clone());
74    Ok(())
75}
76
77pub fn set_status(new_status: &str) -> Result<()> {
78    let mut store = lock_state_store()?;
79    let st = store
80        .registration
81        .get_or_insert_with(RegistrationState::default);
82    st.status = new_status.to_string();
83    Ok(())
84}
85
86pub fn touch_healthcheck_now() -> Result<()> {
87    let mut store = lock_state_store()?;
88    let st = store
89        .registration
90        .get_or_insert_with(RegistrationState::default);
91    st.last_healthcheck = Some(Utc::now());
92    Ok(())
93}
94
95pub struct LockGuard;
96
97impl LockGuard {
98    pub fn try_acquire(stale_after: Duration) -> std::io::Result<Option<Self>> {
99        let mut state = lock_lock_store()?;
100        let now = Instant::now();
101
102        if let Some(acquired_at) = state.acquired_at {
103            if now.duration_since(acquired_at) <= stale_after {
104                return Ok(None);
105            }
106        }
107
108        state.acquired_at = Some(now);
109        state.owner_pid = Some(std::process::id());
110
111        Ok(Some(Self))
112    }
113}
114
115impl Drop for LockGuard {
116    fn drop(&mut self) {
117        if let Ok(mut state) = lock_lock_store() {
118            state.acquired_at = None;
119            state.owner_pid = None;
120        }
121    }
122}
123
124#[derive(Default)]
125struct LockState {
126    acquired_at: Option<Instant>,
127    owner_pid: Option<u32>,
128}
129
130static LOCK_STATE: OnceLock<Mutex<LockState>> = OnceLock::new();
131
132fn lock_store() -> &'static Mutex<LockState> {
133    LOCK_STATE.get_or_init(|| Mutex::new(LockState::default()))
134}
135
136fn lock_lock_store() -> io::Result<MutexGuard<'static, LockState>> {
137    lock_store()
138        .lock()
139        .map_err(|_| io::Error::other("registration lock store poisoned"))
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    #[test]
147    fn default_state_is_disconnected() {
148        let st = RegistrationState::default();
149        assert_eq!(st.status, STATUS_DISCONNECTED);
150        assert!(st.last_healthcheck.is_none());
151    }
152
153    #[test]
154    fn write_and_read_secret_roundtrip() {
155        clear_node_secret().unwrap();
156
157        write_node_secret("first").unwrap();
158        let got = read_node_secret().unwrap();
159        assert_eq!(got.as_deref(), Some("first"));
160
161        write_node_secret("second").unwrap();
162        let got2 = read_node_secret().unwrap();
163        assert_eq!(got2.as_deref(), Some("second"));
164
165        clear_node_secret().unwrap();
166        assert!(read_node_secret().unwrap().is_none());
167    }
168}
posemesh_node_registration/state.rs

posemesh_node_registration/
state.rs
