Skip to main content

byteflow/scheduler/
registry.rs

1//! Named registry: `name → CapId` (address), never `name → FlowId`.
2//!
3//! Entries are swept when the **target** flow exits ([`Registry::unregister_flow`]).
4//! A revoked Cap cannot be registered; lookup after exit returns `None`.
5
6use std::collections::HashMap;
7
8use crate::bytecode::CapId;
9use super::error::{LifecycleError, RuntimeError};
10use super::process::FlowId;
11use super::sync_lock;
12
13/// Registry key. Interned as `Box<str>` so lookups do not allocate a `String`
14/// on the happy path when the caller already has a `&str`.
15#[derive(Debug, Clone, PartialEq, Eq, Hash)]
16pub struct RegistryName(Box<str>);
17
18impl RegistryName {
19    pub fn new(name: impl Into<Box<str>>) -> Self {
20        Self(name.into())
21    }
22
23    pub fn as_str(&self) -> &str {
24        &self.0
25    }
26}
27
28impl From<&str> for RegistryName {
29    fn from(name: &str) -> Self {
30        Self(name.into())
31    }
32}
33
34struct Entry {
35    cap: CapId,
36    flow: FlowId,
37}
38
39pub struct Registry {
40    by_name: HashMap<RegistryName, Entry>,
41    by_flow: HashMap<FlowId, Vec<RegistryName>>,
42}
43
44impl Registry {
45    pub fn new() -> Self {
46        Self {
47            by_name: HashMap::new(),
48            by_flow: HashMap::new(),
49        }
50    }
51
52    pub fn register(
53        &mut self,
54        name: RegistryName,
55        cap: CapId,
56        flow: FlowId,
57    ) -> Result<(), LifecycleError> {
58        if name.as_str().is_empty() {
59            return Err(LifecycleError::EmptyName);
60        }
61        if self.by_name.contains_key(&name) {
62            return Err(LifecycleError::AlreadyRegistered);
63        }
64        self.by_flow
65            .entry(flow)
66            .or_default()
67            .push(name.clone());
68        self.by_name.insert(name, Entry { cap, flow });
69        Ok(())
70    }
71
72    pub fn whereis(&self, name: &str) -> Option<CapId> {
73        self.by_name.get(&RegistryName::from(name)).map(|e| e.cap)
74    }
75
76    pub fn unregister(&mut self, name: &str) -> bool {
77        let key = RegistryName::from(name);
78        match self.by_name.remove(&key) {
79            Some(entry) => {
80                if let Some(names) = self.by_flow.get_mut(&entry.flow) {
81                    names.retain(|n| n != &key);
82                    if names.is_empty() {
83                        self.by_flow.remove(&entry.flow);
84                    }
85                }
86                true
87            }
88            None => false,
89        }
90    }
91
92    /// Drop every name that pointed at `flow` (called from finalize).
93    pub fn unregister_flow(&mut self, flow: FlowId) {
94        if let Some(names) = self.by_flow.remove(&flow) {
95            for name in names {
96                self.by_name.remove(&name);
97            }
98        }
99    }
100}
101
102impl Default for Registry {
103    fn default() -> Self {
104        Self::new()
105    }
106}
107
108pub struct RegistryStore {
109    inner: std::sync::Mutex<Registry>,
110}
111
112impl RegistryStore {
113    pub fn new() -> Self {
114        Self {
115            inner: std::sync::Mutex::new(Registry::new()),
116        }
117    }
118
119    pub fn register(
120        &self,
121        name: RegistryName,
122        cap: CapId,
123        flow: FlowId,
124    ) -> Result<Result<(), LifecycleError>, RuntimeError> {
125        Ok(sync_lock::lock(&self.inner, "RegistryStore::register")?.register(name, cap, flow))
126    }
127
128    pub fn whereis(&self, name: &str) -> Result<Option<CapId>, RuntimeError> {
129        Ok(sync_lock::lock(&self.inner, "RegistryStore::whereis")?.whereis(name))
130    }
131
132    pub fn unregister(&self, name: &str) -> Result<bool, RuntimeError> {
133        Ok(sync_lock::lock(&self.inner, "RegistryStore::unregister")?.unregister(name))
134    }
135
136    pub fn unregister_flow(&self, flow: FlowId) -> Result<(), RuntimeError> {
137        sync_lock::lock(&self.inner, "RegistryStore::unregister_flow")?.unregister_flow(flow);
138        Ok(())
139    }
140}
141
142impl Default for RegistryStore {
143    fn default() -> Self {
144        Self::new()
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151    use crate::bytecode::CapId;
152    use crate::scheduler::process::next_flow_id;
153
154    #[test]
155    fn unregister_flow_clears_names() {
156        let mut reg = Registry::new();
157        let flow = next_flow_id();
158        let cap = CapId::from_raw(42);
159        assert!(reg.register(RegistryName::from("svc"), cap, flow).is_ok());
160        assert_eq!(reg.whereis("svc"), Some(cap));
161        reg.unregister_flow(flow);
162        assert_eq!(reg.whereis("svc"), None);
163    }
164
165    #[test]
166    fn duplicate_name_is_rejected() {
167        let mut reg = Registry::new();
168        let flow = next_flow_id();
169        assert!(reg
170            .register(RegistryName::from("svc"), CapId::from_raw(1), flow)
171            .is_ok());
172        assert_eq!(
173            reg.register(RegistryName::from("svc"), CapId::from_raw(2), flow),
174            Err(LifecycleError::AlreadyRegistered)
175        );
176    }
177
178    #[test]
179    fn empty_name_is_rejected() {
180        let mut reg = Registry::new();
181        let flow = next_flow_id();
182        assert_eq!(
183            reg.register(RegistryName::from(""), CapId::from_raw(1), flow),
184            Err(LifecycleError::EmptyName)
185        );
186    }
187}