1use crate::handler_trait::Handler;
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::collections::HashMap;
10use std::fmt;
11use std::sync::{Arc, RwLock};
12
13#[derive(Debug, Clone)]
15pub struct RegistryError {
16 message: String,
18}
19
20impl RegistryError {
21 fn lock_poisoned() -> Self {
23 Self {
24 message: "Failed to acquire lock on registry: lock was poisoned due to a previous panic".to_string(),
25 }
26 }
27}
28
29impl fmt::Display for RegistryError {
30 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31 write!(f, "{}", self.message)
32 }
33}
34
35impl std::error::Error for RegistryError {}
36
37#[derive(Debug, Clone, Serialize, Deserialize)]
39pub struct MethodExample {
40 pub name: String,
42 pub description: Option<String>,
44 pub params: Value,
46 pub result: Value,
48}
49
50#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct MethodMetadata {
53 pub name: String,
55 pub description: Option<String>,
57 pub params_schema: Option<Value>,
59 pub result_schema: Option<Value>,
61 pub error_schema: Option<Value>,
63 pub examples: Vec<MethodExample>,
65 pub deprecated: bool,
67 pub tags: Vec<String>,
69}
70
71impl MethodMetadata {
72 pub fn new(name: impl Into<String>) -> Self {
78 Self {
79 name: name.into(),
80 description: None,
81 params_schema: None,
82 result_schema: None,
83 error_schema: None,
84 examples: Vec::new(),
85 deprecated: false,
86 tags: Vec::new(),
87 }
88 }
89
90 pub fn with_description(mut self, description: impl Into<String>) -> Self {
92 self.description = Some(description.into());
93 self
94 }
95
96 pub fn with_params_schema(mut self, schema: Value) -> Self {
98 self.params_schema = Some(schema);
99 self
100 }
101
102 pub fn with_result_schema(mut self, schema: Value) -> Self {
104 self.result_schema = Some(schema);
105 self
106 }
107
108 pub fn with_error_schema(mut self, schema: Value) -> Self {
110 self.error_schema = Some(schema);
111 self
112 }
113
114 pub fn with_example(mut self, example: MethodExample) -> Self {
116 self.examples.push(example);
117 self
118 }
119
120 pub fn mark_deprecated(mut self) -> Self {
122 self.deprecated = true;
123 self
124 }
125
126 pub fn with_tag(mut self, tag: impl Into<String>) -> Self {
128 self.tags.push(tag.into());
129 self
130 }
131}
132
133type MethodEntry = (Arc<dyn Handler>, MethodMetadata);
135
136type MethodStorage = Arc<RwLock<HashMap<String, MethodEntry>>>;
138
139type MethodListEntry = (String, Arc<dyn Handler>, MethodMetadata);
141
142pub struct JsonRpcMethodRegistry {
168 methods: MethodStorage,
170}
171
172impl JsonRpcMethodRegistry {
173 pub fn new() -> Self {
175 Self {
176 methods: Arc::new(RwLock::new(HashMap::new())),
177 }
178 }
179
180 pub fn register(
212 &self,
213 name: impl Into<String>,
214 handler: Arc<dyn Handler>,
215 metadata: MethodMetadata,
216 ) -> Result<(), RegistryError> {
217 let name = name.into();
218 let mut methods = self.methods.write().map_err(|_| RegistryError::lock_poisoned())?;
219 methods.insert(name, (handler, metadata));
220 Ok(())
221 }
222
223 pub fn get(&self, name: &str) -> Result<Option<Arc<dyn Handler>>, RegistryError> {
236 let methods = self.methods.read().map_err(|_| RegistryError::lock_poisoned())?;
237 Ok(methods.get(name).map(|(handler, _)| Arc::clone(handler)))
238 }
239
240 pub fn get_metadata(&self, name: &str) -> Result<Option<MethodMetadata>, RegistryError> {
253 let methods = self.methods.read().map_err(|_| RegistryError::lock_poisoned())?;
254 Ok(methods.get(name).map(|(_, metadata)| metadata.clone()))
255 }
256
257 pub fn get_with_metadata(&self, name: &str) -> Result<Option<MethodEntry>, RegistryError> {
270 let methods = self.methods.read().map_err(|_| RegistryError::lock_poisoned())?;
271 Ok(methods
272 .get(name)
273 .map(|(handler, metadata)| (Arc::clone(handler), metadata.clone())))
274 }
275
276 pub fn list_methods(&self) -> Result<Vec<String>, RegistryError> {
283 let methods = self.methods.read().map_err(|_| RegistryError::lock_poisoned())?;
284 let mut names: Vec<String> = methods.keys().cloned().collect();
285 names.sort();
286 Ok(names)
287 }
288
289 pub fn contains(&self, name: &str) -> Result<bool, RegistryError> {
300 let methods = self.methods.read().map_err(|_| RegistryError::lock_poisoned())?;
301 Ok(methods.contains_key(name))
302 }
303
304 pub fn len(&self) -> Result<usize, RegistryError> {
311 let methods = self.methods.read().map_err(|_| RegistryError::lock_poisoned())?;
312 Ok(methods.len())
313 }
314
315 pub fn is_empty(&self) -> Result<bool, RegistryError> {
322 let methods = self.methods.read().map_err(|_| RegistryError::lock_poisoned())?;
323 Ok(methods.is_empty())
324 }
325
326 pub fn remove(&self, name: &str) -> Result<bool, RegistryError> {
339 let mut methods = self.methods.write().map_err(|_| RegistryError::lock_poisoned())?;
340 Ok(methods.remove(name).is_some())
341 }
342
343 pub fn clear(&self) -> Result<(), RegistryError> {
349 let mut methods = self.methods.write().map_err(|_| RegistryError::lock_poisoned())?;
350 methods.clear();
351 Ok(())
352 }
353
354 pub fn list_all(&self) -> Result<Vec<MethodListEntry>, RegistryError> {
361 let methods = self.methods.read().map_err(|_| RegistryError::lock_poisoned())?;
362 Ok(methods
363 .iter()
364 .map(|(name, (handler, metadata))| (name.clone(), Arc::clone(handler), metadata.clone()))
365 .collect())
366 }
367}
368
369impl Default for JsonRpcMethodRegistry {
370 fn default() -> Self {
371 Self::new()
372 }
373}
374
375impl Clone for JsonRpcMethodRegistry {
376 fn clone(&self) -> Self {
377 Self {
378 methods: Arc::clone(&self.methods),
379 }
380 }
381}
382
383#[cfg(test)]
384mod tests {
385 use super::*;
386 use crate::handler_trait::{HandlerResult, RequestData};
387 use axum::body::Body;
388 use axum::http::Request;
389 use std::panic;
390
391 struct MockHandler;
393
394 impl Handler for MockHandler {
395 fn call(
396 &self,
397 _request: Request<Body>,
398 _request_data: RequestData,
399 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = HandlerResult> + Send + '_>> {
400 Box::pin(async { Err((axum::http::StatusCode::OK, "mock".to_string())) })
401 }
402 }
403
404 fn create_test_registry() -> JsonRpcMethodRegistry {
405 JsonRpcMethodRegistry::new()
406 }
407
408 fn create_mock_handler() -> Arc<dyn Handler> {
409 Arc::new(MockHandler)
410 }
411
412 #[test]
413 fn test_new_registry_is_empty() {
414 let registry = create_test_registry();
415 assert!(registry.is_empty().unwrap());
416 assert_eq!(registry.len().unwrap(), 0);
417 assert!(registry.list_methods().unwrap().is_empty());
418 }
419
420 #[test]
421 fn test_register_and_get_method() {
422 let registry = create_test_registry();
423 let handler = create_mock_handler();
424 let metadata = MethodMetadata::new("test_method");
425
426 registry
427 .register("test_method", handler.clone(), metadata.clone())
428 .unwrap();
429
430 assert!(!registry.is_empty().unwrap());
431 assert_eq!(registry.len().unwrap(), 1);
432 assert!(registry.contains("test_method").unwrap());
433 assert!(registry.get("test_method").unwrap().is_some());
434 assert_eq!(
435 registry.get_metadata("test_method").unwrap().unwrap().name,
436 "test_method"
437 );
438 }
439
440 #[test]
441 fn test_get_nonexistent_method() {
442 let registry = create_test_registry();
443 assert!(registry.get("nonexistent").unwrap().is_none());
444 assert!(registry.get_metadata("nonexistent").unwrap().is_none());
445 }
446
447 #[test]
448 fn test_list_methods_returns_sorted() {
449 let registry = create_test_registry();
450 let handler = create_mock_handler();
451
452 registry
453 .register("zebra", handler.clone(), MethodMetadata::new("zebra"))
454 .unwrap();
455 registry
456 .register("apple", handler.clone(), MethodMetadata::new("apple"))
457 .unwrap();
458 registry
459 .register("banana", handler.clone(), MethodMetadata::new("banana"))
460 .unwrap();
461
462 let methods = registry.list_methods().unwrap();
463 assert_eq!(methods, vec!["apple", "banana", "zebra"]);
464 }
465
466 #[test]
467 fn test_register_overwrites_existing() {
468 let registry = create_test_registry();
469 let handler1 = create_mock_handler();
470 let handler2 = create_mock_handler();
471
472 registry
473 .register("method", handler1, MethodMetadata::new("method"))
474 .unwrap();
475 assert_eq!(registry.len().unwrap(), 1);
476
477 registry
478 .register("method", handler2, MethodMetadata::new("method"))
479 .unwrap();
480 assert_eq!(registry.len().unwrap(), 1);
481 }
482
483 #[test]
484 fn test_remove_method() {
485 let registry = create_test_registry();
486 let handler = create_mock_handler();
487
488 registry
489 .register("method", handler, MethodMetadata::new("method"))
490 .unwrap();
491 assert_eq!(registry.len().unwrap(), 1);
492
493 let removed = registry.remove("method").unwrap();
494 assert!(removed);
495 assert_eq!(registry.len().unwrap(), 0);
496 assert!(registry.get("method").unwrap().is_none());
497 }
498
499 #[test]
500 fn test_remove_nonexistent_method() {
501 let registry = create_test_registry();
502 let removed = registry.remove("nonexistent").unwrap();
503 assert!(!removed);
504 }
505
506 #[test]
507 fn test_clear_registry() {
508 let registry = create_test_registry();
509 let handler = create_mock_handler();
510
511 registry
512 .register("method1", handler.clone(), MethodMetadata::new("method1"))
513 .unwrap();
514 registry
515 .register("method2", handler.clone(), MethodMetadata::new("method2"))
516 .unwrap();
517 registry
518 .register("method3", handler.clone(), MethodMetadata::new("method3"))
519 .unwrap();
520
521 assert_eq!(registry.len().unwrap(), 3);
522 registry.clear().unwrap();
523 assert_eq!(registry.len().unwrap(), 0);
524 assert!(registry.is_empty().unwrap());
525 }
526
527 #[test]
528 fn test_get_with_metadata() {
529 let registry = create_test_registry();
530 let handler = create_mock_handler();
531 let metadata = MethodMetadata::new("method").with_description("Test method");
532
533 registry.register("method", handler.clone(), metadata).unwrap();
534
535 let result = registry.get_with_metadata("method").unwrap();
536 assert!(result.is_some());
537
538 let (_retrieved_handler, retrieved_metadata) = result.unwrap();
539 assert_eq!(retrieved_metadata.name, "method");
540 assert_eq!(retrieved_metadata.description, Some("Test method".to_string()));
541 }
542
543 #[test]
544 fn test_list_all() {
545 let registry = create_test_registry();
546 let handler = create_mock_handler();
547
548 registry
549 .register("add", handler.clone(), MethodMetadata::new("add"))
550 .unwrap();
551 registry
552 .register("subtract", handler.clone(), MethodMetadata::new("subtract"))
553 .unwrap();
554
555 let all = registry.list_all().unwrap();
556 assert_eq!(all.len(), 2);
557
558 let names: Vec<String> = all.iter().map(|(name, _, _)| name.clone()).collect();
559 assert!(names.contains(&"add".to_string()));
560 assert!(names.contains(&"subtract".to_string()));
561 }
562
563 #[test]
564 fn test_clone_shares_registry() {
565 let registry1 = create_test_registry();
566 let handler = create_mock_handler();
567
568 registry1
569 .register("method", handler, MethodMetadata::new("method"))
570 .unwrap();
571
572 let registry2 = registry1.clone();
573 assert_eq!(registry2.len().unwrap(), 1);
574 assert!(registry2.contains("method").unwrap());
575 }
576
577 #[test]
578 fn test_metadata_builder_pattern() {
579 let metadata = MethodMetadata::new("test")
580 .with_description("Test description")
581 .with_tag("math")
582 .with_tag("utility")
583 .mark_deprecated();
584
585 assert_eq!(metadata.name, "test");
586 assert_eq!(metadata.description, Some("Test description".to_string()));
587 assert_eq!(metadata.tags, vec!["math", "utility"]);
588 assert!(metadata.deprecated);
589 }
590
591 #[test]
592 fn test_metadata_with_schemas() {
593 let params_schema = serde_json::json!({
594 "type": "object",
595 "properties": {
596 "x": { "type": "number" },
597 "y": { "type": "number" }
598 }
599 });
600
601 let result_schema = serde_json::json!({
602 "type": "number"
603 });
604
605 let metadata = MethodMetadata::new("add")
606 .with_params_schema(params_schema.clone())
607 .with_result_schema(result_schema.clone());
608
609 assert_eq!(metadata.params_schema, Some(params_schema));
610 assert_eq!(metadata.result_schema, Some(result_schema));
611 }
612
613 #[test]
614 fn test_metadata_with_examples() {
615 let example = MethodExample {
616 name: "example1".to_string(),
617 description: Some("Test example".to_string()),
618 params: serde_json::json!({"x": 1, "y": 2}),
619 result: serde_json::json!(3),
620 };
621
622 let metadata = MethodMetadata::new("add").with_example(example.clone());
623
624 assert_eq!(metadata.examples.len(), 1);
625 assert_eq!(metadata.examples[0].name, "example1");
626 assert_eq!(metadata.examples[0].description, Some("Test example".to_string()));
627 }
628
629 #[test]
630 fn test_registry_errors_on_poisoned_lock() {
631 let registry = create_test_registry();
632 let _ = panic::catch_unwind(|| {
633 let _guard = registry.methods.write().expect("lock write");
634 panic!("poison the lock");
635 });
636
637 let handler = create_mock_handler();
638 let err = registry
639 .register("method", handler, MethodMetadata::new("method"))
640 .expect_err("poisoned lock should error");
641 assert!(err.to_string().contains("lock was poisoned"));
642
643 match registry.get("method") {
644 Err(err) => assert!(err.to_string().contains("lock was poisoned")),
645 Ok(_) => panic!("expected poisoned lock error"),
646 }
647 }
648}