1use crate::error::PluginError;
13use crate::main_engine::{PluginInfo, PluginState};
14use crate::service::ServiceDesc;
15use ek_event::{Event, HandlerId};
16use std::any::Any;
17use std::collections::{HashMap, VecDeque};
18use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering};
19use std::sync::{Arc, Mutex, RwLock};
20use std::thread::JoinHandle;
21
22use crossbeam_channel::Sender;
23
24#[derive(Default)]
26pub(crate) struct EntryStats {
27 pub spilled: AtomicU64,
29 pub panics: AtomicU64,
31 pub pending: AtomicU64,
33 pub peak: AtomicU64,
35 pub done: AtomicBool,
37}
38
39pub(crate) struct PluginEntry {
41 pub name: String,
42 pub plugin: Arc<dyn crate::plugin::Plugin>,
44 pub any: Arc<dyn Any + Send + Sync>,
46 pub state: PluginStateCell,
47 pub soft_cap: u64,
49 pub deps: Vec<String>,
50 pub mailbox: Mutex<Option<Sender<Event>>>,
52 pub forwarder_ids: Mutex<Vec<HandlerId>>,
54 pub worker: Mutex<Option<JoinHandle<()>>>,
56 pub stats: Arc<EntryStats>,
57}
58
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61pub enum PluginStateCode {
62 Ready = 0,
63 Running = 1,
64 Stopping = 2,
65 Exited = 3,
66 Quarantined = 4,
67}
68
69#[derive(Clone)]
71pub(crate) struct PluginStateCell(Arc<AtomicU8>);
72
73impl PluginStateCell {
74 pub fn new(s: PluginStateCode) -> Self {
75 Self(Arc::new(AtomicU8::new(s as u8)))
76 }
77 pub fn get(&self) -> PluginState {
78 PluginState::from_code(self.0.load(Ordering::Acquire))
79 }
80 pub(crate) fn cas(&self, from: PluginStateCode, to: PluginStateCode) -> bool {
81 self.0
82 .compare_exchange(from as u8, to as u8, Ordering::AcqRel, Ordering::Acquire)
83 .is_ok()
84 }
85 pub(crate) fn finish(&self) -> bool {
87 self.cas(PluginStateCode::Running, PluginStateCode::Exited)
88 || self.cas(PluginStateCode::Stopping, PluginStateCode::Exited)
89 }
90}
91
92pub(crate) struct RegistryInner {
94 pub entries: HashMap<String, Arc<PluginEntry>>,
95 pub order: Vec<String>,
97}
98
99type SvcCacheEntry = (u64, Arc<Vec<ServiceDesc>>);
101
102pub(crate) struct Registry {
103 pub state_of: RwLock<RegistryInner>,
104 pub(crate) svc_index: RwLock<HashMap<String, Vec<String>>>,
106 revision: AtomicU64,
108 svc_cache: Mutex<HashMap<String, SvcCacheEntry>>,
110}
111
112impl Registry {
113 pub fn new() -> Self {
114 Self {
115 state_of: RwLock::new(RegistryInner {
116 entries: HashMap::new(),
117 order: Vec::new(),
118 }),
119 svc_index: RwLock::new(HashMap::new()),
120 revision: AtomicU64::new(1),
121 svc_cache: Mutex::new(HashMap::new()),
122 }
123 }
124
125 pub fn revision(&self) -> u64 {
126 self.revision.load(Ordering::Acquire)
127 }
128
129 pub fn bump(&self) {
130 self.revision.fetch_add(1, Ordering::AcqRel);
131 }
132
133 pub fn index_put(&self, name: &str, ids: &[String]) {
135 let mut w = self.svc_index.write().unwrap();
136 for list in w.values_mut() {
137 list.retain(|n| n != name);
138 }
139 w.retain(|_, l| !l.is_empty());
140 for id in ids {
141 let entry = w.entry(id.clone()).or_default();
142 if !entry.iter().any(|n| n == name) {
143 entry.push(name.to_string());
144 }
145 }
146 }
147
148 pub fn index_remove(&self, name: &str) {
150 let mut w = self.svc_index.write().unwrap();
151 for list in w.values_mut() {
152 list.retain(|n| n != name);
153 }
154 w.retain(|_, l| !l.is_empty());
155 }
156
157 pub fn index_providers(&self, svc: &str) -> Vec<String> {
159 self.svc_index
160 .read()
161 .unwrap()
162 .get(svc)
163 .cloned()
164 .unwrap_or_default()
165 }
166
167 pub(crate) fn publish(&self, name: &str, list: Vec<ServiceDesc>) -> String {
170 let Some(entry) = self.get(name) else {
171 return String::new();
172 };
173 let mut ids: Vec<String> = Vec::with_capacity(list.len());
174 for d in &list {
175 if !ids.contains(&d.id) {
176 ids.push(d.id.clone());
177 }
178 }
179 self.index_put(name, &ids);
180 self.invalidate(name);
181 entry.plugin.version().into_owned()
182 }
183
184 pub fn services_of(&self, name: &str) -> Option<Arc<Vec<ServiceDesc>>> {
186 let entry = self.get(name)?;
187 let revision = self.revision();
188 if let Some((rev, list)) = self.svc_cache.lock().unwrap().get(name) {
189 if *rev == revision {
190 return Some(Arc::clone(list));
191 }
192 }
193 let list = Arc::new(entry.plugin.services());
194 self.svc_cache
195 .lock()
196 .unwrap()
197 .insert(name.to_string(), (revision, Arc::clone(&list)));
198 Some(list)
199 }
200
201 pub(crate) fn info_of(&self, name: &str) -> Option<PluginInfo> {
203 self.get(name).map(|e| PluginInfo {
204 name: e.name.clone(),
205 version: e.plugin.version().into_owned(),
206 state: e.state.get(),
207 deps: e.deps.clone(),
208 })
209 }
210
211 pub fn get(&self, name: &str) -> Option<Arc<PluginEntry>> {
212 self.state_of.read().unwrap().entries.get(name).cloned()
213 }
214
215 pub fn remove(&self, name: &str) -> Option<Arc<PluginEntry>> {
216 let mut w = self.state_of.write().unwrap();
217 w.entries.remove(name).inspect(|_| {
218 if let Some(i) = w.order.iter().position(|n| n == name) {
219 w.order.remove(i);
220 }
221 })
222 }
223
224 pub fn names(&self) -> Vec<String> {
225 self.state_of.read().unwrap().order.clone()
226 }
227
228 pub fn any_depends_on(&self, target: &str) -> Vec<String> {
229 let r = self.state_of.read().unwrap();
230 r.order
231 .iter()
232 .filter_map(|n| r.entries.get(n))
233 .filter(|e| e.deps.iter().any(|d| d == target))
234 .map(|e| e.name.clone())
235 .collect()
236 }
237
238 pub fn clear(&self) {
239 let mut w = self.state_of.write().unwrap();
240 w.entries.clear();
241 w.order.clear();
242 drop(w);
243 self.svc_index.write().unwrap().clear();
244 self.svc_cache.lock().unwrap().clear();
245 self.bump();
246 }
247
248 pub fn invalidate(&self, name: &str) {
250 self.svc_cache.lock().unwrap().remove(name);
251 self.bump();
252 }
253
254 pub fn topo_order(&self) -> Result<Vec<Arc<PluginEntry>>, PluginError> {
257 let r = self.state_of.read().unwrap();
258 let mut indeg: HashMap<String, usize> = HashMap::new();
259 let mut dependents: HashMap<String, Vec<String>> = HashMap::new();
260 for name in &r.order {
261 indeg.entry(name.clone()).or_insert(0);
262 }
263 for name in &r.order {
264 let e = &r.entries[name];
265 for d in &e.deps {
266 if !r.entries.contains_key(d) {
267 return Err(PluginError::UnknownDep(d.clone()));
268 }
269 *indeg.entry(name.clone()).or_insert(0) += 1;
270 dependents.entry(d.clone()).or_default().push(name.clone());
271 }
272 }
273 let mut queue: VecDeque<String> =
274 r.order.iter().filter(|n| indeg[*n] == 0).cloned().collect();
275 let mut out = Vec::with_capacity(r.order.len());
276 while let Some(n) = queue.pop_front() {
277 let e = Arc::clone(&r.entries[&n]);
278 out.push(e);
279 if let Some(deps_list) = dependents.get(&n) {
280 for dep_name in deps_list {
281 let d = indeg.get_mut(dep_name).unwrap();
282 *d -= 1;
283 if *d == 0 {
284 queue.push_back(dep_name.clone());
285 }
286 }
287 }
288 }
289 if out.len() < r.order.len() {
290 let cycle: Vec<String> = r.order.iter().filter(|n| indeg[*n] > 0).cloned().collect();
291 return Err(PluginError::DepCycle(cycle));
292 }
293 Ok(out)
294 }
295}