zus_rpc_client/

client.rs

//! ZUS RPC Client with Connection Pooling Support
//!
//! Provides high-level RPC client with optional connection pooling
//! for high-throughput scenarios like gateways.

use {async_trait::async_trait, bytes::Bytes, tracing::instrument};

use zus_common::{ConnectionPoolConfig, Result};

use crate::endpoint::ZooLBRpcEndpoint;

/// Abstract ZUS RPC Client (replacing Java's AbstractZusRpcClient)
#[async_trait]
pub trait ZusRpcClient: Send + Sync {
  /// Synchronous RPC call
  async fn sync_call(&self, method: &str, request: Bytes, timeout_ms: u64) -> Result<Bytes>;

  /// Asynchronous RPC call (fire and forget)
  async fn notify_call(&self, method: &str, request: Bytes) -> Result<()>;
}

/// Base RPC Client implementation
pub struct BaseRpcClient {
  endpoint: ZooLBRpcEndpoint,
  #[allow(dead_code)]
  load_balance_enabled: bool,
}

impl BaseRpcClient {
  /// Create client with default configuration (no pooling)
  ///
  /// # Example
  /// ```ignore
  /// let client = BaseRpcClient::new("zns://zoo1:2181/services/myservice").await?;
  /// ```
  pub async fn new(server_path: &str) -> Result<Self> {
    let endpoint = ZooLBRpcEndpoint::new(server_path).await?;

    Ok(Self {
      endpoint,
      load_balance_enabled: true,
    })
  }

  /// Create client with connection pooling enabled
  ///
  /// # Example
  /// ```ignore
  /// let client = BaseRpcClient::new_with_pooling(
  ///     "zns://zoo1:2181/services/myservice",
  ///     ConnectionPoolConfig::for_gateway()
  /// ).await?;
  /// ```
  pub async fn new_with_pooling(server_path: &str, pool_config: ConnectionPoolConfig) -> Result<Self> {
    let endpoint = ZooLBRpcEndpoint::new_with_pooling(server_path, pool_config).await?;

    Ok(Self {
      endpoint,
      load_balance_enabled: true,
    })
  }

  /// Create a builder for more advanced configuration
  pub fn builder() -> BaseRpcClientBuilder {
    BaseRpcClientBuilder::new()
  }

  /// Create client optimized for gateway scenarios
  ///
  /// Uses connection pooling with sensible defaults for high-throughput:
  /// - min_connections: 2
  /// - max_connections: 10
  ///
  /// # Example
  /// ```ignore
  /// let client = BaseRpcClient::for_gateway("zns://zoo1:2181/services/backend").await?;
  /// ```
  pub async fn for_gateway(server_path: &str) -> Result<Self> {
    Self::new_with_pooling(server_path, ConnectionPoolConfig::for_gateway()).await
  }

  pub fn enable_load_balance(&mut self) {
    self.load_balance_enabled = true;
  }

  pub fn disable_load_balance(&mut self) {
    self.load_balance_enabled = false;
  }

  /// Check if connection pooling is enabled
  pub fn is_pooling_enabled(&self) -> bool {
    self.endpoint.is_pooling_enabled()
  }
}

#[async_trait]
impl ZusRpcClient for BaseRpcClient {
  #[instrument(name = "client_sync_call", skip(self, request), fields(method = %method, timeout_ms = timeout_ms))]
  async fn sync_call(&self, method: &str, request: Bytes, timeout_ms: u64) -> Result<Bytes> {
    let server = self.endpoint.select_one_server().await?;
    let method_bytes = Bytes::copy_from_slice(method.as_bytes());
    server.sync_call(method_bytes, request, timeout_ms).await
  }

  #[instrument(name = "client_notify_call", skip(self, request), fields(method = %method))]
  async fn notify_call(&self, method: &str, request: Bytes) -> Result<()> {
    let server = self.endpoint.select_one_server().await?;
    let method_bytes = Bytes::copy_from_slice(method.as_bytes());
    server.notify_call(method_bytes, request).await
  }
}

/// Builder for BaseRpcClient with advanced configuration
///
/// # Example
/// ```ignore
/// let client = BaseRpcClient::builder()
///     .address("zns://zoo1:2181/services/myservice")
///     .pool_config(ConnectionPoolConfig {
///         min_connections: 4,
///         max_connections: 20,
///         ..Default::default()
///     })
///     .build()
///     .await?;
/// ```
pub struct BaseRpcClientBuilder {
  address: Option<String>,
  pool_config: Option<ConnectionPoolConfig>,
  load_balance_enabled: bool,
}

impl BaseRpcClientBuilder {
  pub fn new() -> Self {
    Self {
      address: None,
      pool_config: None,
      load_balance_enabled: true,
    }
  }

  /// Set the server address (required)
  ///
  /// Supports:
  /// - `zns://host:port/service/path` for service discovery
  /// - `tcp://host:port` for direct connection
  pub fn address(mut self, address: &str) -> Self {
    self.address = Some(address.to_string());
    self
  }

  /// Enable connection pooling with the specified configuration
  pub fn pool_config(mut self, config: ConnectionPoolConfig) -> Self {
    self.pool_config = Some(config);
    self
  }

  /// Use gateway-optimized pooling configuration
  pub fn for_gateway(mut self) -> Self {
    self.pool_config = Some(ConnectionPoolConfig::for_gateway());
    self
  }

  /// Use custom pool size
  pub fn with_pool_size(mut self, min: usize, max: usize) -> Self {
    self.pool_config = Some(ConnectionPoolConfig::with_pool_size(min, max));
    self
  }

  /// Enable/disable load balancing (default: enabled)
  pub fn load_balance(mut self, enabled: bool) -> Self {
    self.load_balance_enabled = enabled;
    self
  }

  /// Build the client
  pub async fn build(self) -> Result<BaseRpcClient> {
    let address = self
      .address
      .ok_or_else(|| zus_common::ZusError::Protocol("Address is required".to_string()))?;

    let endpoint = ZooLBRpcEndpoint::with_pool_config(&address, self.pool_config).await?;

    Ok(BaseRpcClient {
      endpoint,
      load_balance_enabled: self.load_balance_enabled,
    })
  }
}

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