arora_websocket/
registry.rs

1//! Registry for slots and methods.
2//!
3//! Provides thread-safe storage for available slots and invocable methods.
4
5use tokio::sync::RwLock;
6
7use std::collections::HashMap;
8use std::sync::Arc;
9
10use crate::method::{InvokeResult, MethodInfo};
11use crate::slot::SlotInfo;
12use arora_types::value::Value;
13
14/// Trait for method handlers in the registry.
15///
16/// Implement this trait to create custom method handlers.
17pub trait RegistryMethodHandler: Send + Sync {
18    /// Handle a method invocation.
19    fn invoke(&self, args: HashMap<String, Value>) -> InvokeResult;
20}
21
22/// Function-based method handler.
23impl<F> RegistryMethodHandler for F
24where
25    F: Fn(HashMap<String, Value>) -> InvokeResult + Send + Sync,
26{
27    fn invoke(&self, args: HashMap<String, Value>) -> InvokeResult {
28        self(args)
29    }
30}
31
32/// Registry for slots and methods.
33///
34/// This is the core state container for the WebSocket server.
35/// It stores available slots and registered methods.
36pub struct Registry {
37    slots: RwLock<Vec<SlotInfo>>,
38    methods: RwLock<HashMap<String, MethodInfo>>,
39    handlers: RwLock<HashMap<String, Arc<dyn RegistryMethodHandler>>>,
40}
41
42impl Default for Registry {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl Registry {
49    /// Create a new empty registry.
50    pub fn new() -> Self {
51        Self {
52            slots: RwLock::new(Vec::new()),
53            methods: RwLock::new(HashMap::new()),
54            handlers: RwLock::new(HashMap::new()),
55        }
56    }
57
58    /// Set the available slots.
59    pub async fn set_slots(&self, slots: Vec<SlotInfo>) {
60        *self.slots.write().await = slots;
61    }
62
63    /// Get all registered slots.
64    pub async fn get_slots(&self) -> Vec<SlotInfo> {
65        self.slots.read().await.clone()
66    }
67
68    /// Get slots filtered by path prefix.
69    pub async fn get_slots_filtered(&self, prefix: Option<&str>) -> Vec<SlotInfo> {
70        let slots = self.slots.read().await;
71        match prefix {
72            Some(prefix) => {
73                let prefix = prefix.trim_end_matches('/');
74                slots
75                    .iter()
76                    .filter(|n| {
77                        n.path.starts_with(prefix) || n.path.starts_with(&format!("{}/", prefix))
78                    })
79                    .cloned()
80                    .collect()
81            }
82            None => slots.clone(),
83        }
84    }
85
86    /// Get input slots (slots with kind == "input").
87    pub async fn get_input_paths(&self) -> Vec<String> {
88        self.slots
89            .read()
90            .await
91            .iter()
92            .filter(|n| n.kind.as_deref() == Some("input"))
93            .map(|n| n.path.clone())
94            .collect()
95    }
96
97    /// Register a method with its handler.
98    pub async fn register_method<H>(&self, info: MethodInfo, handler: H)
99    where
100        H: RegistryMethodHandler + 'static,
101    {
102        let path = info.path.clone();
103        self.methods.write().await.insert(path.clone(), info);
104        self.handlers.write().await.insert(path, Arc::new(handler));
105    }
106
107    /// Register a method using a closure.
108    pub async fn register_method_fn<F>(&self, info: MethodInfo, handler: F)
109    where
110        F: Fn(HashMap<String, Value>) -> InvokeResult + Send + Sync + 'static,
111    {
112        self.register_method(info, handler).await;
113    }
114
115    /// Get all registered methods.
116    pub async fn get_methods(&self) -> Vec<MethodInfo> {
117        self.methods.read().await.values().cloned().collect()
118    }
119
120    /// Get methods filtered by path prefix.
121    pub async fn get_methods_filtered(&self, prefix: Option<&str>) -> Vec<MethodInfo> {
122        let methods = self.methods.read().await;
123        match prefix {
124            Some(prefix) => {
125                let prefix = prefix.trim_end_matches('/');
126                methods
127                    .values()
128                    .filter(|m| {
129                        m.path.starts_with(prefix) || m.path.starts_with(&format!("{}/", prefix))
130                    })
131                    .cloned()
132                    .collect()
133            }
134            None => methods.values().cloned().collect(),
135        }
136    }
137
138    /// Invoke a method by path.
139    pub async fn invoke_method(&self, path: &str, args: HashMap<String, Value>) -> InvokeResult {
140        let handlers = self.handlers.read().await;
141        match handlers.get(path) {
142            Some(handler) => {
143                let handler = handler.clone();
144                drop(handlers); // Release lock before invoking
145                handler.invoke(args)
146            }
147            None => InvokeResult::err(format!("Method not found: {}", path)),
148        }
149    }
150
151    /// Check if a method exists.
152    pub async fn has_method(&self, path: &str) -> bool {
153        self.handlers.read().await.contains_key(path)
154    }
155}
arora_websocket/registry.rs

arora_websocket/
registry.rs
