use crate::error::PluginError;
use crate::main_engine::{PluginInfo, PluginState};
use crate::service::ServiceDesc;
use ek_event::{Event, HandlerId};
use std::any::Any;
use std::collections::{HashMap, VecDeque};
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use std::thread::JoinHandle;
use crossbeam_channel::Sender;
#[derive(Default)]
pub(crate) struct EntryStats {
pub spilled: AtomicU64,
pub panics: AtomicU64,
pub done: AtomicBool,
}
pub(crate) struct PluginEntry {
pub name: String,
pub plugin: Arc<dyn crate::plugin::Plugin>,
pub any: Arc<dyn Any + Send + Sync>,
pub state: PluginStateCell,
pub mailbox_cap: usize,
pub deps: Vec<String>,
pub mailbox: Mutex<Option<Sender<Event>>>,
pub forwarder_ids: Mutex<Vec<HandlerId>>,
pub worker: Mutex<Option<JoinHandle<()>>>,
pub stats: Arc<EntryStats>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PluginStateCode {
Ready = 0,
Running = 1,
Stopping = 2,
Exited = 3,
Quarantined = 4,
}
#[derive(Clone)]
pub(crate) struct PluginStateCell(Arc<AtomicU8>);
impl PluginStateCell {
pub fn new(s: PluginStateCode) -> Self {
Self(Arc::new(AtomicU8::new(s as u8)))
}
pub fn get(&self) -> PluginState {
PluginState::from_code(self.0.load(Ordering::Acquire))
}
pub(crate) fn cas(&self, from: PluginStateCode, to: PluginStateCode) -> bool {
self.0
.compare_exchange(from as u8, to as u8, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
}
pub(crate) fn finish(&self) -> bool {
self.cas(PluginStateCode::Running, PluginStateCode::Exited)
|| self.cas(PluginStateCode::Stopping, PluginStateCode::Exited)
}
}
pub(crate) struct RegistryInner {
pub entries: HashMap<String, Arc<PluginEntry>>,
pub order: Vec<String>,
}
type SvcCacheEntry = (u64, Arc<Vec<ServiceDesc>>);
pub(crate) struct Registry {
pub state_of: RwLock<RegistryInner>,
pub(crate) svc_index: RwLock<HashMap<String, Vec<String>>>,
revision: AtomicU64,
svc_cache: Mutex<HashMap<String, SvcCacheEntry>>,
}
impl Registry {
pub fn new() -> Self {
Self {
state_of: RwLock::new(RegistryInner {
entries: HashMap::new(),
order: Vec::new(),
}),
svc_index: RwLock::new(HashMap::new()),
revision: AtomicU64::new(1),
svc_cache: Mutex::new(HashMap::new()),
}
}
pub fn revision(&self) -> u64 {
self.revision.load(Ordering::Acquire)
}
pub fn bump(&self) {
self.revision.fetch_add(1, Ordering::AcqRel);
}
pub fn index_put(&self, name: &str, ids: &[String]) {
let mut w = self.svc_index.write().unwrap();
for list in w.values_mut() {
list.retain(|n| n != name);
}
w.retain(|_, l| !l.is_empty());
for id in ids {
let entry = w.entry(id.clone()).or_default();
if !entry.iter().any(|n| n == name) {
entry.push(name.to_string());
}
}
}
pub fn index_remove(&self, name: &str) {
let mut w = self.svc_index.write().unwrap();
for list in w.values_mut() {
list.retain(|n| n != name);
}
w.retain(|_, l| !l.is_empty());
}
pub fn index_providers(&self, svc: &str) -> Vec<String> {
self.svc_index
.read()
.unwrap()
.get(svc)
.cloned()
.unwrap_or_default()
}
pub(crate) fn publish(&self, name: &str, list: Vec<ServiceDesc>) -> String {
let Some(entry) = self.get(name) else {
return String::new();
};
let mut ids: Vec<String> = Vec::with_capacity(list.len());
for d in &list {
if !ids.contains(&d.id) {
ids.push(d.id.clone());
}
}
self.index_put(name, &ids);
self.invalidate(name);
entry.plugin.version().into_owned()
}
pub fn services_of(&self, name: &str) -> Option<Arc<Vec<ServiceDesc>>> {
let entry = self.get(name)?;
let revision = self.revision();
if let Some((rev, list)) = self.svc_cache.lock().unwrap().get(name) {
if *rev == revision {
return Some(Arc::clone(list));
}
}
let list = Arc::new(entry.plugin.services());
self.svc_cache
.lock()
.unwrap()
.insert(name.to_string(), (revision, Arc::clone(&list)));
Some(list)
}
pub(crate) fn info_of(&self, name: &str) -> Option<PluginInfo> {
self.get(name).map(|e| PluginInfo {
name: e.name.clone(),
version: e.plugin.version().into_owned(),
state: e.state.get(),
deps: e.deps.clone(),
})
}
pub fn get(&self, name: &str) -> Option<Arc<PluginEntry>> {
self.state_of.read().unwrap().entries.get(name).cloned()
}
pub fn remove(&self, name: &str) -> Option<Arc<PluginEntry>> {
let mut w = self.state_of.write().unwrap();
w.entries.remove(name).inspect(|_| {
if let Some(i) = w.order.iter().position(|n| n == name) {
w.order.remove(i);
}
})
}
pub fn names(&self) -> Vec<String> {
self.state_of.read().unwrap().order.clone()
}
pub fn any_depends_on(&self, target: &str) -> Vec<String> {
let r = self.state_of.read().unwrap();
r.order
.iter()
.filter_map(|n| r.entries.get(n))
.filter(|e| e.deps.iter().any(|d| d == target))
.map(|e| e.name.clone())
.collect()
}
pub fn clear(&self) {
let mut w = self.state_of.write().unwrap();
w.entries.clear();
w.order.clear();
drop(w);
self.svc_index.write().unwrap().clear();
self.svc_cache.lock().unwrap().clear();
self.bump();
}
pub fn invalidate(&self, name: &str) {
self.svc_cache.lock().unwrap().remove(name);
self.bump();
}
pub fn topo_order(&self) -> Result<Vec<Arc<PluginEntry>>, PluginError> {
let r = self.state_of.read().unwrap();
let mut indeg: HashMap<String, usize> = HashMap::new();
let mut dependents: HashMap<String, Vec<String>> = HashMap::new();
for name in &r.order {
indeg.entry(name.clone()).or_insert(0);
}
for name in &r.order {
let e = &r.entries[name];
for d in &e.deps {
if !r.entries.contains_key(d) {
return Err(PluginError::UnknownDep(d.clone()));
}
*indeg.entry(name.clone()).or_insert(0) += 1;
dependents.entry(d.clone()).or_default().push(name.clone());
}
}
let mut queue: VecDeque<String> =
r.order.iter().filter(|n| indeg[*n] == 0).cloned().collect();
let mut out = Vec::with_capacity(r.order.len());
while let Some(n) = queue.pop_front() {
let e = Arc::clone(&r.entries[&n]);
out.push(e);
if let Some(deps_list) = dependents.get(&n) {
for dep_name in deps_list {
let d = indeg.get_mut(dep_name).unwrap();
*d -= 1;
if *d == 0 {
queue.push_back(dep_name.clone());
}
}
}
}
if out.len() < r.order.len() {
let cycle: Vec<String> = r.order.iter().filter(|n| indeg[*n] > 0).cloned().collect();
return Err(PluginError::DepCycle(cycle));
}
Ok(out)
}
}