arora_websocket/
registry.rs

1//! Registry for keys and methods.
2//!
3//! Provides thread-safe storage for the advertised keys and invocable methods.
4
5use tokio::sync::RwLock;
6
7use std::collections::HashMap;
8use std::sync::Arc;
9
10use crate::key::KeyInfo;
11use crate::method::{InvokeResult, MethodInfo};
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 keys and methods.
33///
34/// This is the core state container for the WebSocket server.
35/// It stores the advertised keys and registered methods.
36pub struct Registry {
37    keys: RwLock<Vec<KeyInfo>>,
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            keys: RwLock::new(Vec::new()),
53            methods: RwLock::new(HashMap::new()),
54            handlers: RwLock::new(HashMap::new()),
55        }
56    }
57
58    /// Set the advertised keys.
59    pub async fn set_keys(&self, keys: Vec<KeyInfo>) {
60        *self.keys.write().await = keys;
61    }
62
63    /// Get all advertised keys.
64    pub async fn get_keys(&self) -> Vec<KeyInfo> {
65        self.keys.read().await.clone()
66    }
67
68    /// Get keys filtered by path prefix.
69    pub async fn get_keys_filtered(&self, prefix: Option<&str>) -> Vec<KeyInfo> {
70        let keys = self.keys.read().await;
71        match prefix {
72            Some(prefix) => {
73                let prefix = prefix.trim_end_matches('/');
74                keys.iter()
75                    .filter(|n| {
76                        n.path.starts_with(prefix) || n.path.starts_with(&format!("{}/", prefix))
77                    })
78                    .cloned()
79                    .collect()
80            }
81            None => keys.clone(),
82        }
83    }
84
85    /// Get input key paths (keys with kind == "input").
86    pub async fn get_input_paths(&self) -> Vec<String> {
87        self.keys
88            .read()
89            .await
90            .iter()
91            .filter(|n| n.kind.as_deref() == Some("input"))
92            .map(|n| n.path.clone())
93            .collect()
94    }
95
96    /// Register a method with its handler.
97    pub async fn register_method<H>(&self, info: MethodInfo, handler: H)
98    where
99        H: RegistryMethodHandler + 'static,
100    {
101        let path = info.path.clone();
102        self.methods.write().await.insert(path.clone(), info);
103        self.handlers.write().await.insert(path, Arc::new(handler));
104    }
105
106    /// Register a method using a closure.
107    pub async fn register_method_fn<F>(&self, info: MethodInfo, handler: F)
108    where
109        F: Fn(HashMap<String, Value>) -> InvokeResult + Send + Sync + 'static,
110    {
111        self.register_method(info, handler).await;
112    }
113
114    /// Get all registered methods.
115    pub async fn get_methods(&self) -> Vec<MethodInfo> {
116        self.methods.read().await.values().cloned().collect()
117    }
118
119    /// Get methods filtered by path prefix.
120    pub async fn get_methods_filtered(&self, prefix: Option<&str>) -> Vec<MethodInfo> {
121        let methods = self.methods.read().await;
122        match prefix {
123            Some(prefix) => {
124                let prefix = prefix.trim_end_matches('/');
125                methods
126                    .values()
127                    .filter(|m| {
128                        m.path.starts_with(prefix) || m.path.starts_with(&format!("{}/", prefix))
129                    })
130                    .cloned()
131                    .collect()
132            }
133            None => methods.values().cloned().collect(),
134        }
135    }
136
137    /// Invoke a method by path.
138    pub async fn invoke_method(&self, path: &str, args: HashMap<String, Value>) -> InvokeResult {
139        let handlers = self.handlers.read().await;
140        match handlers.get(path) {
141            Some(handler) => {
142                let handler = handler.clone();
143                drop(handlers); // Release lock before invoking
144                handler.invoke(args)
145            }
146            None => InvokeResult::err(format!("Method not found: {}", path)),
147        }
148    }
149
150    /// Check if a method exists.
151    pub async fn has_method(&self, path: &str) -> bool {
152        self.handlers.read().await.contains_key(path)
153    }
154}
arora_websocket/registry.rs

arora_websocket/
registry.rs
