starpc 0.41.2 → 0.43.1

package/srpc/mux.rs

@@ -0,0 +1,353 @@
+//! Service multiplexer for routing RPC calls.
+//!
+//! The Mux provides a registry for RPC handlers and routes incoming calls
+//! to the appropriate handler based on service and method IDs.
+//!
+//! This implementation matches the Go Mux behavior, supporting:
+//! - Service-level handler registration
+//! - Method-level routing within services
+//! - Fallback invokers for unmatched calls
+//! - Empty service ID handling (searches all services)
+
+use async_trait::async_trait;
+use std::collections::HashMap;
+use std::sync::{Arc, RwLock};
+
+use crate::error::{Error, Result};
+use crate::handler::Handler;
+use crate::invoker::Invoker;
+use crate::stream::Stream;
+
+/// A shared handler reference.
+pub type ArcHandler = Arc<dyn Handler>;
+
+/// Method handlers map: method_id -> handler.
+type MethodHandlers = HashMap<String, ArcHandler>;
+
+/// Service multiplexer for routing RPC calls to handlers.
+///
+/// The Mux maintains a registry of handlers and routes incoming RPC calls
+/// to the appropriate handler based on service and method IDs.
+///
+/// # Example
+///
+/// ```ignore
+/// use starpc::{Mux, Handler};
+///
+/// let mux = Mux::new();
+/// mux.register(Arc::new(MyServiceHandler))?;
+///
+/// // The mux can now route calls to MyServiceHandler
+/// ```
+pub struct Mux {
+    /// Map of service ID -> method ID -> handler.
+    /// This matches the Go implementation's structure.
+    services: RwLock<HashMap<String, MethodHandlers>>,
+
+    /// Fallback invokers to try if no handler is found.
+    fallbacks: RwLock<Vec<Arc<dyn Invoker>>>,
+}
+
+impl Default for Mux {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Mux {
+    /// Creates a new empty Mux.
+    pub fn new() -> Self {
+        Self {
+            services: RwLock::new(HashMap::new()),
+            fallbacks: RwLock::new(Vec::new()),
+        }
+    }
+
+    /// Creates a new Mux with fallback invokers.
+    ///
+    /// Fallback invokers are called in order when no handler matches
+    /// the requested service/method.
+    pub fn with_fallbacks(fallbacks: Vec<Arc<dyn Invoker>>) -> Self {
+        Self {
+            services: RwLock::new(HashMap::new()),
+            fallbacks: RwLock::new(fallbacks),
+        }
+    }
+
+    /// Registers a handler for its service.
+    ///
+    /// The handler is registered for all methods it declares via `method_ids()`.
+    ///
+    /// # Errors
+    ///
+    /// Returns `Error::EmptyServiceId` if the handler's service ID is empty.
+    pub fn register(&self, handler: ArcHandler) -> Result<()> {
+        let service_id = handler.service_id();
+        if service_id.is_empty() {
+            return Err(Error::EmptyServiceId);
+        }
+
+        let method_ids = handler.method_ids();
+
+        let mut services = self.services.write().unwrap();
+        let service_methods = services
+            .entry(service_id.to_string())
+            .or_insert_with(HashMap::new);
+
+        for method_id in method_ids {
+            if !method_id.is_empty() {
+                service_methods.insert(method_id.to_string(), handler.clone());
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Adds a fallback invoker.
+    ///
+    /// Fallback invokers are tried in order when no handler matches.
+    pub fn add_fallback(&self, invoker: Arc<dyn Invoker>) {
+        self.fallbacks.write().unwrap().push(invoker);
+    }
+
+    /// Checks if a service is registered.
+    pub fn has_service(&self, service_id: &str) -> bool {
+        if service_id.is_empty() {
+            return false;
+        }
+
+        let services = self.services.read().unwrap();
+        services
+            .get(service_id)
+            .map(|methods| !methods.is_empty())
+            .unwrap_or(false)
+    }
+
+    /// Checks if a service method is registered.
+    pub fn has_service_method(&self, service_id: &str, method_id: &str) -> bool {
+        if service_id.is_empty() || method_id.is_empty() {
+            return false;
+        }
+
+        let services = self.services.read().unwrap();
+        services
+            .get(service_id)
+            .and_then(|methods| methods.get(method_id))
+            .is_some()
+    }
+
+    /// Gets a handler by service ID and method ID.
+    fn get_handler(&self, service_id: &str, method_id: &str) -> Option<ArcHandler> {
+        let services = self.services.read().unwrap();
+        services
+            .get(service_id)
+            .and_then(|methods| methods.get(method_id).cloned())
+    }
+
+    /// Finds a handler that implements the given method (searching all services).
+    ///
+    /// This is used when service_id is empty.
+    fn find_handler_for_method(&self, method_id: &str) -> Option<ArcHandler> {
+        let services = self.services.read().unwrap();
+        for methods in services.values() {
+            if let Some(handler) = methods.get(method_id) {
+                return Some(handler.clone());
+            }
+        }
+        None
+    }
+}
+
+/// QueryableInvoker allows checking if methods are implemented without invoking.
+///
+/// This matches the Go interface of the same name.
+pub trait QueryableInvoker {
+    /// Checks if the service is registered.
+    fn has_service(&self, service_id: &str) -> bool;
+
+    /// Checks if the service method is registered.
+    fn has_service_method(&self, service_id: &str, method_id: &str) -> bool;
+}
+
+impl QueryableInvoker for Mux {
+    fn has_service(&self, service_id: &str) -> bool {
+        Mux::has_service(self, service_id)
+    }
+
+    fn has_service_method(&self, service_id: &str, method_id: &str) -> bool {
+        Mux::has_service_method(self, service_id, method_id)
+    }
+}
+
+#[async_trait]
+impl Invoker for Mux {
+    async fn invoke_method(
+        &self,
+        service_id: &str,
+        method_id: &str,
+        stream: Box<dyn Stream>,
+    ) -> (bool, Result<()>) {
+        // Look up the handler
+        let handler = if service_id.is_empty() {
+            // If service_id is empty, search all services for the method
+            self.find_handler_for_method(method_id)
+        } else {
+            self.get_handler(service_id, method_id)
+        };
+
+        if let Some(h) = handler {
+            // Invoke the handler
+            return h.invoke_method(service_id, method_id, stream).await;
+        }
+
+        // Try the first fallback invoker.
+        // Only the first fallback gets a chance because the stream is consumed by the call.
+        // This matches Go behavior where the stream cannot be reused after an invocation attempt.
+        let maybe_fallback = self.fallbacks.read().unwrap().first().cloned();
+        if let Some(fallback) = maybe_fallback {
+            return fallback.invoke_method(service_id, method_id, stream).await;
+        }
+
+        (false, Err(Error::Unimplemented))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::stream::Context;
+
+    struct TestHandler {
+        service_id: &'static str,
+        method_ids: &'static [&'static str],
+    }
+
+    #[async_trait]
+    impl Invoker for TestHandler {
+        async fn invoke_method(
+            &self,
+            _service_id: &str,
+            _method_id: &str,
+            _stream: Box<dyn Stream>,
+        ) -> (bool, Result<()>) {
+            (true, Ok(()))
+        }
+    }
+
+    impl Handler for TestHandler {
+        fn service_id(&self) -> &'static str {
+            self.service_id
+        }
+
+        fn method_ids(&self) -> &'static [&'static str] {
+            self.method_ids
+        }
+    }
+
+    #[test]
+    fn test_mux_register() {
+        let mux = Mux::new();
+        let handler = Arc::new(TestHandler {
+            service_id: "test.Service",
+            method_ids: &["Method1", "Method2"],
+        });
+
+        mux.register(handler).unwrap();
+
+        assert!(mux.has_service("test.Service"));
+        assert!(mux.has_service_method("test.Service", "Method1"));
+        assert!(mux.has_service_method("test.Service", "Method2"));
+        assert!(!mux.has_service_method("test.Service", "Method3"));
+        assert!(!mux.has_service("other.Service"));
+    }
+
+    #[test]
+    fn test_mux_register_empty_service_id() {
+        let mux = Mux::new();
+        let handler = Arc::new(TestHandler {
+            service_id: "",
+            method_ids: &["Method1"],
+        });
+
+        let result = mux.register(handler);
+        assert!(matches!(result, Err(Error::EmptyServiceId)));
+    }
+
+    #[test]
+    fn test_mux_has_service_empty_id() {
+        let mux = Mux::new();
+        assert!(!mux.has_service(""));
+    }
+
+    #[test]
+    fn test_mux_has_service_method_empty_ids() {
+        let mux = Mux::new();
+        let handler = Arc::new(TestHandler {
+            service_id: "test.Service",
+            method_ids: &["Method1"],
+        });
+        mux.register(handler).unwrap();
+
+        assert!(!mux.has_service_method("", "Method1"));
+        assert!(!mux.has_service_method("test.Service", ""));
+    }
+
+    #[test]
+    fn test_mux_find_handler_for_method() {
+        let mux = Mux::new();
+        let handler = Arc::new(TestHandler {
+            service_id: "test.Service",
+            method_ids: &["UniqueMethod"],
+        });
+        mux.register(handler).unwrap();
+
+        // Should find handler when searching all services
+        let found = mux.find_handler_for_method("UniqueMethod");
+        assert!(found.is_some());
+
+        let not_found = mux.find_handler_for_method("NonExistent");
+        assert!(not_found.is_none());
+    }
+
+    #[tokio::test]
+    async fn test_mux_invoke_with_empty_service_id() {
+        let mux = Mux::new();
+        let handler = Arc::new(TestHandler {
+            service_id: "test.Service",
+            method_ids: &["TestMethod"],
+        });
+        mux.register(handler).unwrap();
+
+        // Create a mock stream
+        struct MockStream;
+
+        #[async_trait]
+        impl Stream for MockStream {
+            fn context(&self) -> &Context {
+                static CTX: std::sync::OnceLock<Context> = std::sync::OnceLock::new();
+                CTX.get_or_init(Context::new)
+            }
+
+            async fn send_bytes(&self, _data: bytes::Bytes) -> Result<()> {
+                Ok(())
+            }
+
+            async fn recv_bytes(&self) -> Result<bytes::Bytes> {
+                Err(Error::StreamClosed)
+            }
+
+            async fn close_send(&self) -> Result<()> {
+                Ok(())
+            }
+
+            async fn close(&self) -> Result<()> {
+                Ok(())
+            }
+        }
+
+        // Should find the method even with empty service ID
+        let (handled, result) = mux.invoke_method("", "TestMethod", Box::new(MockStream)).await;
+        assert!(handled);
+        assert!(result.is_ok());
+    }
+}

package/srpc/packet.rs

@@ -0,0 +1,334 @@
+//! Packet utility functions and validation.
+//!
+//! This module provides helper functions for creating and validating
+//! starpc protocol packets.
+
+use crate::error::{Error, Result};
+use crate::proto::{packet::Body, CallData, CallStart, Packet};
+use crate::transport::encode_optional_data;
+use bytes::Bytes;
+
+/// Creates a new CallStart packet.
+///
+/// # Arguments
+/// * `service` - The service ID
+/// * `method` - The method ID
+/// * `data` - Optional initial data (None = no data, Some(empty) = empty data)
+pub fn new_call_start(
+    service: impl Into<String>,
+    method: impl Into<String>,
+    data: Option<Bytes>,
+) -> Packet {
+    let (data_bytes, data_is_zero) = encode_optional_data(data);
+
+    Packet {
+        body: Some(Body::CallStart(CallStart {
+            rpc_service: service.into(),
+            rpc_method: method.into(),
+            data: data_bytes,
+            data_is_zero,
+        })),
+    }
+}
+
+/// Creates a new CallData packet with data.
+///
+/// # Arguments
+/// * `data` - The data to send (can be empty)
+pub fn new_call_data(data: Vec<u8>) -> Packet {
+    let data_is_zero = data.is_empty();
+    Packet {
+        body: Some(Body::CallData(CallData {
+            data,
+            data_is_zero,
+            complete: false,
+            error: String::new(),
+        })),
+    }
+}
+
+/// Creates a new CallData packet with optional data and flags.
+///
+/// # Arguments
+/// * `data` - Optional data bytes
+/// * `data_is_zero` - True if empty data should be sent
+/// * `complete` - True if this completes the stream
+/// * `error` - Optional error message
+pub fn new_call_data_full(
+    data: Option<Bytes>,
+    complete: bool,
+    error: Option<String>,
+) -> Packet {
+    let (data_bytes, data_is_zero) = encode_optional_data(data);
+
+    Packet {
+        body: Some(Body::CallData(CallData {
+            data: data_bytes,
+            data_is_zero,
+            complete: complete || error.is_some(),
+            error: error.unwrap_or_default(),
+        })),
+    }
+}
+
+/// Creates a new CallData packet indicating completion.
+pub fn new_call_complete() -> Packet {
+    Packet {
+        body: Some(Body::CallData(CallData {
+            data: vec![],
+            data_is_zero: false,
+            complete: true,
+            error: String::new(),
+        })),
+    }
+}
+
+/// Creates a new CallData packet with an error.
+///
+/// # Arguments
+/// * `error` - The error message
+pub fn new_call_error(error: impl Into<String>) -> Packet {
+    Packet {
+        body: Some(Body::CallData(CallData {
+            data: vec![],
+            data_is_zero: false,
+            complete: true,
+            error: error.into(),
+        })),
+    }
+}
+
+/// Creates a new CallCancel packet.
+pub fn new_call_cancel() -> Packet {
+    Packet {
+        body: Some(Body::CallCancel(true)),
+    }
+}
+
+/// Extracts the body type name for debugging.
+pub fn body_type_name(packet: &Packet) -> &'static str {
+    match &packet.body {
+        Some(Body::CallStart(_)) => "CallStart",
+        Some(Body::CallData(_)) => "CallData",
+        Some(Body::CallCancel(_)) => "CallCancel",
+        None => "Empty",
+    }
+}
+
+/// Packet validation trait.
+///
+/// Provides validation methods matching the Go implementation's
+/// `Packet.Validate()`, `CallStart.Validate()`, and `CallData.Validate()`.
+pub trait Validate {
+    /// Validates the packet/message, returning an error if invalid.
+    fn validate(&self) -> Result<()>;
+}
+
+impl Validate for Packet {
+    fn validate(&self) -> Result<()> {
+        match &self.body {
+            Some(Body::CallStart(cs)) => cs.validate(),
+            Some(Body::CallData(cd)) => cd.validate(),
+            Some(Body::CallCancel(_)) => Ok(()),
+            None => Err(Error::EmptyPacket),
+        }
+    }
+}
+
+impl Validate for CallStart {
+    fn validate(&self) -> Result<()> {
+        if self.rpc_method.is_empty() {
+            return Err(Error::EmptyMethodId);
+        }
+        if self.rpc_service.is_empty() {
+            return Err(Error::EmptyServiceId);
+        }
+        Ok(())
+    }
+}
+
+impl Validate for CallData {
+    fn validate(&self) -> Result<()> {
+        // A CallData packet must have at least one of:
+        // - Non-empty data
+        // - data_is_zero flag set (indicating intentionally empty data)
+        // - complete flag set
+        // - error message
+        if self.data.is_empty()
+            && !self.data_is_zero
+            && !self.complete
+            && self.error.is_empty()
+        {
+            return Err(Error::EmptyPacket);
+        }
+        Ok(())
+    }
+}
+
+/// Extension trait for Packet to check body type.
+impl Packet {
+    /// Returns true if this is a CallStart packet.
+    pub fn is_call_start(&self) -> bool {
+        matches!(&self.body, Some(Body::CallStart(_)))
+    }
+
+    /// Returns true if this is a CallData packet.
+    pub fn is_call_data(&self) -> bool {
+        matches!(&self.body, Some(Body::CallData(_)))
+    }
+
+    /// Returns true if this is a CallCancel packet.
+    pub fn is_call_cancel(&self) -> bool {
+        matches!(&self.body, Some(Body::CallCancel(true)))
+    }
+
+    /// Extracts the CallStart body, if present.
+    pub fn into_call_start(self) -> Option<CallStart> {
+        match self.body {
+            Some(Body::CallStart(cs)) => Some(cs),
+            _ => None,
+        }
+    }
+
+    /// Extracts the CallData body, if present.
+    pub fn into_call_data(self) -> Option<CallData> {
+        match self.body {
+            Some(Body::CallData(cd)) => Some(cd),
+            _ => None,
+        }
+    }
+
+    /// Returns true if this packet indicates completion (CallData with complete=true or error).
+    pub fn is_complete(&self) -> bool {
+        match &self.body {
+            Some(Body::CallData(cd)) => cd.complete || !cd.error.is_empty(),
+            Some(Body::CallCancel(true)) => true,
+            _ => false,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_new_call_start() {
+        let pkt = new_call_start("test.Service", "Method", None);
+        let cs = pkt.into_call_start().unwrap();
+        assert_eq!(cs.rpc_service, "test.Service");
+        assert_eq!(cs.rpc_method, "Method");
+        assert!(cs.data.is_empty());
+        assert!(!cs.data_is_zero);
+    }
+
+    #[test]
+    fn test_new_call_start_with_empty_data() {
+        let pkt = new_call_start("svc", "method", Some(Bytes::new()));
+        let cs = pkt.into_call_start().unwrap();
+        assert!(cs.data.is_empty());
+        assert!(cs.data_is_zero);
+    }
+
+    #[test]
+    fn test_new_call_start_with_data() {
+        let pkt = new_call_start("svc", "method", Some(Bytes::from(vec![1, 2, 3])));
+        let cs = pkt.into_call_start().unwrap();
+        assert_eq!(cs.data, vec![1, 2, 3]);
+        assert!(!cs.data_is_zero);
+    }
+
+    #[test]
+    fn test_validate_call_start_valid() {
+        let cs = CallStart {
+            rpc_service: "svc".into(),
+            rpc_method: "method".into(),
+            data: vec![],
+            data_is_zero: false,
+        };
+        assert!(cs.validate().is_ok());
+    }
+
+    #[test]
+    fn test_validate_call_start_empty_method() {
+        let cs = CallStart {
+            rpc_service: "svc".into(),
+            rpc_method: String::new(),
+            data: vec![],
+            data_is_zero: false,
+        };
+        assert!(matches!(cs.validate(), Err(Error::EmptyMethodId)));
+    }
+
+    #[test]
+    fn test_validate_call_start_empty_service() {
+        let cs = CallStart {
+            rpc_service: String::new(),
+            rpc_method: "method".into(),
+            data: vec![],
+            data_is_zero: false,
+        };
+        assert!(matches!(cs.validate(), Err(Error::EmptyServiceId)));
+    }
+
+    #[test]
+    fn test_validate_call_data_valid_with_data() {
+        let cd = CallData {
+            data: vec![1, 2, 3],
+            data_is_zero: false,
+            complete: false,
+            error: String::new(),
+        };
+        assert!(cd.validate().is_ok());
+    }
+
+    #[test]
+    fn test_validate_call_data_valid_with_complete() {
+        let cd = CallData {
+            data: vec![],
+            data_is_zero: false,
+            complete: true,
+            error: String::new(),
+        };
+        assert!(cd.validate().is_ok());
+    }
+
+    #[test]
+    fn test_validate_call_data_valid_with_error() {
+        let cd = CallData {
+            data: vec![],
+            data_is_zero: false,
+            complete: false,
+            error: "some error".into(),
+        };
+        assert!(cd.validate().is_ok());
+    }
+
+    #[test]
+    fn test_validate_call_data_valid_with_zero_data() {
+        let cd = CallData {
+            data: vec![],
+            data_is_zero: true,
+            complete: false,
+            error: String::new(),
+        };
+        assert!(cd.validate().is_ok());
+    }
+
+    #[test]
+    fn test_validate_call_data_invalid_empty() {
+        let cd = CallData {
+            data: vec![],
+            data_is_zero: false,
+            complete: false,
+            error: String::new(),
+        };
+        assert!(matches!(cd.validate(), Err(Error::EmptyPacket)));
+    }
+
+    #[test]
+    fn test_validate_packet_empty() {
+        let pkt = Packet { body: None };
+        assert!(matches!(pkt.validate(), Err(Error::EmptyPacket)));
+    }
+}

package/srpc/proto/mod.rs

@@ -0,0 +1,10 @@
+//! Generated protocol buffer types for starpc.
+//!
+//! This module contains the Packet, CallStart, CallData types generated from
+//! rpcproto.proto. These types define the wire protocol for starpc.
+
+// Include the generated protobuf types.
+include!(concat!(env!("OUT_DIR"), "/srpc.rs"));
+
+// Re-export commonly used items.
+pub use self::packet::Body;