synapse_rpc/

registry.rs

//! Interface and instance registry
//!
//! The registry tracks which interfaces are available and routes
//! incoming RPC requests to the appropriate handlers.

use async_trait::async_trait;
use bytes::Bytes;
use std::{
    collections::{HashMap, HashSet},
    sync::Arc,
};
use synapse_primitives::{InstanceId, InterfaceId, MethodId};
use synapse_proto::{RpcResponse, RpcStatus};
use tokio::sync::RwLock;

/// Handler trait for RPC methods
#[async_trait]
pub trait RpcHandler: Send + Sync {
    async fn handle(&self, request: synapse_proto::RpcRequest) -> RpcResponse;
}

/// Registration info for an interface
#[derive(Debug, Clone)]
pub struct InterfaceRegistration {
    pub interface_id: InterfaceId,
    pub interface_version: u32,
    pub method_ids: HashSet<MethodId>,
    pub method_names: Vec<String>,
    pub instance_id: InstanceId,
    pub service_name: String,
    pub interface_name: String,
}

impl InterfaceRegistration {
    /// Create a new interface registration
    pub fn new(
        interface_name: &str,
        service_name: &str,
        method_names: &[&str],
        instance_id: InstanceId,
    ) -> Self {
        Self {
            interface_id: InterfaceId::from_name(interface_name),
            interface_version: 1_000_000, // 1.0.0
            method_ids: method_names
                .iter()
                .map(|n| MethodId::from_name(n))
                .collect(),
            method_names: method_names.iter().map(|n| n.to_string()).collect(),
            instance_id,
            service_name: service_name.to_string(),
            interface_name: interface_name.to_string(),
        }
    }

    /// Set the interface version
    pub fn with_version(mut self, major: u16, minor: u16, patch: u16) -> Self {
        self.interface_version =
            (major as u32) * 1_000_000 + (minor as u32) * 1_000 + (patch as u32);
        self
    }
}

/// Registry for interfaces and their handlers
pub struct InterfaceRegistry {
    interfaces: Arc<RwLock<HashMap<InterfaceId, InterfaceEntry>>>,
}

struct InterfaceEntry {
    registration: InterfaceRegistration,
    handler: Arc<dyn RpcHandler>,
}

impl InterfaceRegistry {
    pub fn new() -> Self {
        Self {
            interfaces: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Register an interface with its handler
    pub async fn register(
        &self,
        registration: InterfaceRegistration,
        handler: Arc<dyn RpcHandler>,
    ) -> Result<(), String> {
        let mut interfaces = self.interfaces.write().await;

        if interfaces.contains_key(&registration.interface_id) {
            return Err(format!(
                "Interface {} already registered",
                registration.interface_name
            ));
        }

        interfaces.insert(
            registration.interface_id,
            InterfaceEntry {
                registration,
                handler,
            },
        );

        Ok(())
    }

    /// Deregister an interface
    pub async fn deregister(&self, interface_id: InterfaceId) -> Result<(), String> {
        let mut interfaces = self.interfaces.write().await;

        if interfaces.remove(&interface_id).is_none() {
            return Err(format!("Interface {:?} not found", interface_id));
        }

        Ok(())
    }

    /// Route a request to the appropriate handler
    pub async fn route(&self, request: synapse_proto::RpcRequest) -> RpcResponse {
        let interfaces = self.interfaces.read().await;

        let interface_id = InterfaceId::from_raw(request.interface_id);
        let method_id = MethodId::from_raw(request.method_id);

        match interfaces.get(&interface_id) {
            Some(entry) => {
                // Check if method is supported
                if !entry.registration.method_ids.contains(&method_id) {
                    return RpcResponse {
                        status: RpcStatus::MethodNotFound as i32,
                        payload: Bytes::new(),
                        error: Some(synapse_proto::RpcError {
                            code: 1004,
                            message: format!("Method {:?} not found", method_id),
                            details: vec![],
                        }),
                        headers: vec![],
                        responded_at_unix_ms: chrono::Utc::now().timestamp_millis(),
                    };
                }

                // Call the handler
                entry.handler.handle(request).await
            }
            None => {
                // Interface not found error
                RpcResponse {
                    status: RpcStatus::InterfaceNotFound as i32,
                    payload: Bytes::new(),
                    error: Some(synapse_proto::RpcError {
                        code: 1003,
                        message: format!("Interface {:?} not found", interface_id),
                        details: vec![],
                    }),
                    headers: vec![],
                    responded_at_unix_ms: chrono::Utc::now().timestamp_millis(),
                }
            }
        }
    }

    /// List all registered interfaces
    pub async fn list_interfaces(&self) -> Vec<InterfaceRegistration> {
        let interfaces = self.interfaces.read().await;
        interfaces
            .values()
            .map(|entry| entry.registration.clone())
            .collect()
    }

    /// Check if an interface is registered
    pub async fn has_interface(&self, interface_id: InterfaceId) -> bool {
        let interfaces = self.interfaces.read().await;
        interfaces.contains_key(&interface_id)
    }
}

impl Default for InterfaceRegistry {
    fn default() -> Self {
        Self::new()
    }
}

/// Simple function-based RPC handler
pub struct FunctionHandler<F>
where
    F: Fn(synapse_proto::RpcRequest) -> futures::future::BoxFuture<'static, RpcResponse>
        + Send
        + Sync,
{
    func: F,
}

impl<F> FunctionHandler<F>
where
    F: Fn(synapse_proto::RpcRequest) -> futures::future::BoxFuture<'static, RpcResponse>
        + Send
        + Sync,
{
    pub fn new(func: F) -> Self {
        Self { func }
    }
}

#[async_trait]
impl<F> RpcHandler for FunctionHandler<F>
where
    F: Fn(synapse_proto::RpcRequest) -> futures::future::BoxFuture<'static, RpcResponse>
        + Send
        + Sync,
{
    async fn handle(&self, request: synapse_proto::RpcRequest) -> RpcResponse {
        (self.func)(request).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_registry_basic() {
        let registry = InterfaceRegistry::new();

        let registration = InterfaceRegistration::new(
            "test.Interface",
            "test-service",
            &["TestMethod"],
            InstanceId::new_random(),
        );

        let interface_id = registration.interface_id;
        let method_id = *registration.method_ids.iter().next().unwrap();

        // Create a simple handler
        let handler = Arc::new(FunctionHandler::new(move |_req| {
            Box::pin(async move {
                RpcResponse {
                    status: RpcStatus::Ok as i32,
                    payload: Bytes::from("test response"),
                    error: None,
                    headers: vec![],
                    responded_at_unix_ms: chrono::Utc::now().timestamp_millis(),
                }
            })
        }));

        // Register interface
        registry.register(registration, handler).await.unwrap();

        // Test routing
        let request = synapse_proto::RpcRequest {
            interface_id: interface_id.into(),
            method_id: method_id.into(),
            headers: vec![],
            payload: Bytes::from("test payload"),
            sent_at_unix_ms: chrono::Utc::now().timestamp_millis(),
        };

        let response = registry.route(request).await;
        assert_eq!(response.status, RpcStatus::Ok as i32);
    }

    #[tokio::test]
    async fn test_interface_not_found() {
        let registry = InterfaceRegistry::new();

        let request = synapse_proto::RpcRequest {
            interface_id: InterfaceId::from_name("nonexistent.Interface").into(),
            method_id: MethodId::from_name("Method").into(),
            headers: vec![],
            payload: Bytes::new(),
            sent_at_unix_ms: 0,
        };

        let response = registry.route(request).await;
        assert_eq!(response.status, RpcStatus::InterfaceNotFound as i32);
    }

    #[tokio::test]
    async fn test_method_not_found() {
        let registry = InterfaceRegistry::new();

        let registration = InterfaceRegistration::new(
            "test.Interface",
            "test-service",
            &["MethodA"],
            InstanceId::new_random(),
        );

        let interface_id = registration.interface_id;

        let handler = Arc::new(FunctionHandler::new(move |_req| {
            Box::pin(async move {
                RpcResponse {
                    status: RpcStatus::Ok as i32,
                    payload: Bytes::new(),
                    error: None,
                    headers: vec![],
                    responded_at_unix_ms: 0,
                }
            })
        }));

        registry.register(registration, handler).await.unwrap();

        // Try calling a method that doesn't exist
        let request = synapse_proto::RpcRequest {
            interface_id: interface_id.into(),
            method_id: MethodId::from_name("MethodB").into(), // Not registered
            headers: vec![],
            payload: Bytes::new(),
            sent_at_unix_ms: 0,
        };

        let response = registry.route(request).await;
        assert_eq!(response.status, RpcStatus::MethodNotFound as i32);
    }
}