byteflow/scheduler/
registry.rs1use std::collections::HashMap;
7
8use super::capability::CapId;
9use super::error::{LifecycleError, RuntimeError};
10use super::process::FlowId;
11use super::sync_lock;
12
13#[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 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::scheduler::capability::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(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(1), flow)
171 .is_ok());
172 assert_eq!(
173 reg.register(RegistryName::from("svc"), CapId(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(1), flow),
184 Err(LifecycleError::EmptyName)
185 );
186 }
187}