batata_client/naming/

naming_service.rs

use std::sync::Arc;
use std::time::Duration;

use parking_lot::RwLock;
use tokio::sync::Notify;
use tokio::task::JoinHandle;
use tracing::{debug, info, warn};

use crate::api::naming::{
    Instance, InstanceRequest, InstanceResponse, QueryServiceResponse,
    ServiceListRequest, ServiceListResponse, ServiceQueryRequest, SubscribeServiceRequest,
    SubscribeServiceResponse,
};
use crate::common::DEFAULT_GROUP;
use crate::error::{BatataError, Result};
use crate::naming::{
    CallbackServiceListener, ServiceChangeEvent, ServiceInfoCache, ServiceListener,
    SubscriberRegistry,
};
use crate::remote::RpcClient;

/// Naming service for service discovery and registration
pub struct NamingService {
    /// RPC client for server communication
    rpc_client: Arc<RpcClient>,

    /// Service information cache
    cache: Arc<ServiceInfoCache>,

    /// Subscriber registry
    subscribers: Arc<SubscriberRegistry>,

    /// Namespace
    namespace: String,

    /// Group name
    group_name: String,

    /// Whether the service is started
    started: Arc<RwLock<bool>>,

    /// Heartbeat task handle
    heartbeat_task: Arc<RwLock<Option<JoinHandle<()>>>>,

    /// Shutdown notify
    shutdown: Arc<Notify>,

    /// Registered instances
    registered_instances: Arc<RwLock<Vec<(String, String, Instance)>>>,
}

impl NamingService {
    /// Create a new NamingService
    pub fn new(rpc_client: Arc<RpcClient>, namespace: &str) -> Self {
        Self {
            rpc_client,
            cache: Arc::new(ServiceInfoCache::new()),
            subscribers: Arc::new(SubscriberRegistry::new()),
            namespace: namespace.to_string(),
            group_name: DEFAULT_GROUP.to_string(),
            started: Arc::new(RwLock::new(false)),
            heartbeat_task: Arc::new(RwLock::new(None)),
            shutdown: Arc::new(Notify::new()),
            registered_instances: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// Set default group name
    pub fn with_group(mut self, group_name: &str) -> Self {
        self.group_name = group_name.to_string();
        self
    }

    /// Start the naming service
    pub async fn start(&self) -> Result<()> {
        if *self.started.read() {
            return Err(BatataError::ClientAlreadyStarted);
        }

        *self.started.write() = true;

        // Start heartbeat task for ephemeral instances
        let instances = self.registered_instances.clone();
        let rpc_client = self.rpc_client.clone();
        let namespace = self.namespace.clone();
        let shutdown = self.shutdown.clone();

        let handle = tokio::spawn(async move {
            Self::heartbeat_loop(instances, rpc_client, namespace, shutdown).await;
        });

        *self.heartbeat_task.write() = Some(handle);

        info!("NamingService started");
        Ok(())
    }

    /// Stop the naming service
    pub async fn stop(&self) {
        *self.started.write() = false;
        self.shutdown.notify_one();

        if let Some(handle) = self.heartbeat_task.write().take() {
            handle.abort();
        }

        // Deregister all instances
        let instances = self.registered_instances.read().clone();
        for (service_name, group_name, instance) in instances {
            if let Err(e) = self.deregister_instance(&service_name, &group_name, instance).await {
                warn!("Failed to deregister instance on shutdown: {}", e);
            }
        }

        info!("NamingService stopped");
    }

    /// Register a service instance
    pub async fn register_instance(
        &self,
        service_name: &str,
        group_name: &str,
        instance: Instance,
    ) -> Result<()> {
        let group_name = if group_name.is_empty() {
            &self.group_name
        } else {
            group_name
        };

        let mut instance = instance;
        instance.generate_instance_id();

        let request =
            InstanceRequest::register(&self.namespace, service_name, group_name, instance.clone());

        let response: InstanceResponse = self.rpc_client.request(&request).await?;

        if !response.response.success {
            return Err(BatataError::server_error(
                response.response.error_code,
                response.response.message,
            ));
        }

        // Track registered instance for heartbeat
        if instance.ephemeral {
            self.registered_instances.write().push((
                service_name.to_string(),
                group_name.to_string(),
                instance,
            ));
        }

        info!(
            "Registered instance: service={}, group={}",
            service_name, group_name
        );

        Ok(())
    }

    /// Register a service instance with simplified parameters
    pub async fn register_instance_simple(
        &self,
        service_name: &str,
        ip: &str,
        port: i32,
    ) -> Result<()> {
        let instance = Instance::new(ip, port);
        self.register_instance(service_name, &self.group_name, instance)
            .await
    }

    /// Deregister a service instance
    pub async fn deregister_instance(
        &self,
        service_name: &str,
        group_name: &str,
        instance: Instance,
    ) -> Result<()> {
        let group_name = if group_name.is_empty() {
            &self.group_name
        } else {
            group_name
        };

        let request =
            InstanceRequest::deregister(&self.namespace, service_name, group_name, instance.clone());

        let response: InstanceResponse = self.rpc_client.request(&request).await?;

        if !response.response.success {
            return Err(BatataError::server_error(
                response.response.error_code,
                response.response.message,
            ));
        }

        // Remove from registered instances
        self.registered_instances
            .write()
            .retain(|(s, g, i)| !(s == service_name && g == group_name && i.key() == instance.key()));

        info!(
            "Deregistered instance: service={}, group={}",
            service_name, group_name
        );

        Ok(())
    }

    /// Deregister a service instance with simplified parameters
    pub async fn deregister_instance_simple(
        &self,
        service_name: &str,
        ip: &str,
        port: i32,
    ) -> Result<()> {
        let instance = Instance::new(ip, port);
        self.deregister_instance(service_name, &self.group_name, instance)
            .await
    }

    /// Update a service instance
    pub async fn update_instance(
        &self,
        service_name: &str,
        group_name: &str,
        instance: Instance,
    ) -> Result<()> {
        let group_name = if group_name.is_empty() {
            &self.group_name
        } else {
            group_name
        };

        let request =
            InstanceRequest::update(&self.namespace, service_name, group_name, instance.clone());

        let response: InstanceResponse = self.rpc_client.request(&request).await?;

        if !response.response.success {
            return Err(BatataError::server_error(
                response.response.error_code,
                response.response.message,
            ));
        }

        // Update in registered instances list
        {
            let mut registered = self.registered_instances.write();
            for (s, g, i) in registered.iter_mut() {
                if s == service_name && g == group_name && i.key() == instance.key() {
                    *i = instance.clone();
                    break;
                }
            }
        }

        info!(
            "Updated instance: service={}, group={}",
            service_name, group_name
        );

        Ok(())
    }

    /// Update a service instance with simplified parameters
    pub async fn update_instance_simple(
        &self,
        service_name: &str,
        ip: &str,
        port: i32,
        weight: f64,
        enabled: bool,
    ) -> Result<()> {
        let instance = Instance::new(ip, port)
            .with_weight(weight)
            .with_enabled(enabled);
        self.update_instance(service_name, &self.group_name, instance)
            .await
    }

    /// Get service information with cluster filtering
    pub async fn get_service(
        &self,
        service_name: &str,
        group_name: &str,
        clusters: &[String],
    ) -> Result<crate::api::naming::Service> {
        let group_name = if group_name.is_empty() {
            &self.group_name
        } else {
            group_name
        };

        let cluster_str = clusters.join(",");
        let request = ServiceQueryRequest::new(&self.namespace, service_name, group_name)
            .with_cluster(&cluster_str);

        let response: QueryServiceResponse = self.rpc_client.request(&request).await?;

        // Update cache
        self.cache.put(&self.namespace, response.service_info.clone());

        Ok(response.service_info)
    }

    /// Get all instances of a service
    pub async fn get_all_instances(
        &self,
        service_name: &str,
        group_name: &str,
    ) -> Result<Vec<Instance>> {
        self.select_instances(service_name, group_name, false).await
    }

    /// Select healthy instances of a service
    pub async fn select_instances(
        &self,
        service_name: &str,
        group_name: &str,
        healthy_only: bool,
    ) -> Result<Vec<Instance>> {
        let group_name = if group_name.is_empty() {
            &self.group_name
        } else {
            group_name
        };

        // Try cache first
        if let Some(service) = self.cache.get(&self.namespace, group_name, service_name) {
            let instances = if healthy_only {
                service
                    .hosts
                    .into_iter()
                    .filter(|i| i.healthy && i.enabled)
                    .collect()
            } else {
                service.hosts
            };
            return Ok(instances);
        }

        // Fetch from server
        let request =
            ServiceQueryRequest::new(&self.namespace, service_name, group_name)
                .with_healthy_only(healthy_only);

        let response: QueryServiceResponse = self.rpc_client.request(&request).await?;

        // Update cache
        self.cache.put(&self.namespace, response.service_info.clone());

        let instances = if healthy_only {
            response
                .service_info
                .hosts
                .into_iter()
                .filter(|i| i.healthy && i.enabled)
                .collect()
        } else {
            response.service_info.hosts
        };

        Ok(instances)
    }

    /// Select one healthy instance randomly
    pub async fn select_one_healthy_instance(
        &self,
        service_name: &str,
        group_name: &str,
    ) -> Result<Instance> {
        let instances = self.select_instances(service_name, group_name, true).await?;

        if instances.is_empty() {
            return Err(BatataError::ServiceNotFound {
                service_name: service_name.to_string(),
                group_name: group_name.to_string(),
                namespace: self.namespace.clone(),
            });
        }

        // Simple random selection
        let index = rand_index(instances.len());
        Ok(instances[index].clone())
    }

    /// Get list of services
    pub async fn get_services_of_server(
        &self,
        group_name: &str,
        page_no: i32,
        page_size: i32,
    ) -> Result<(i32, Vec<String>)> {
        let group_name = if group_name.is_empty() {
            &self.group_name
        } else {
            group_name
        };

        let request = ServiceListRequest::new(&self.namespace, group_name)
            .with_page(page_no, page_size);

        let response: ServiceListResponse = self.rpc_client.request(&request).await?;

        Ok((response.count, response.service_names))
    }

    /// Subscribe to service changes
    pub async fn subscribe<L>(
        &self,
        service_name: &str,
        group_name: &str,
        listener: L,
    ) -> Result<()>
    where
        L: ServiceListener + 'static,
    {
        let group_name = if group_name.is_empty() {
            &self.group_name
        } else {
            group_name
        };

        // Add listener
        self.subscribers.subscribe(
            &self.namespace,
            group_name,
            service_name,
            Arc::new(listener),
        );

        // Send subscribe request
        let request = SubscribeServiceRequest::subscribe(&self.namespace, service_name, group_name);

        let response: SubscribeServiceResponse = self.rpc_client.request(&request).await?;

        // Update cache with initial data
        self.cache.put(&self.namespace, response.service_info.clone());

        info!(
            "Subscribed to service: service={}, group={}",
            service_name, group_name
        );

        Ok(())
    }

    /// Subscribe with callback
    pub async fn subscribe_callback<F>(
        &self,
        service_name: &str,
        group_name: &str,
        callback: F,
    ) -> Result<()>
    where
        F: Fn(ServiceChangeEvent) + Send + Sync + 'static,
    {
        self.subscribe(service_name, group_name, CallbackServiceListener::new(callback))
            .await
    }

    /// Unsubscribe from service changes
    pub async fn unsubscribe(&self, service_name: &str, group_name: &str) -> Result<()> {
        let group_name = if group_name.is_empty() {
            &self.group_name
        } else {
            group_name
        };

        self.subscribers
            .unsubscribe(&self.namespace, group_name, service_name);

        // Send unsubscribe request
        let request = SubscribeServiceRequest::unsubscribe(&self.namespace, service_name, group_name);

        let _response: SubscribeServiceResponse = self.rpc_client.request(&request).await?;

        info!(
            "Unsubscribed from service: service={}, group={}",
            service_name, group_name
        );

        Ok(())
    }

    /// Get server status
    pub async fn get_server_status(&self) -> Result<String> {
        if self.rpc_client.is_connected() {
            Ok("UP".to_string())
        } else {
            Ok("DOWN".to_string())
        }
    }

    /// Background heartbeat loop for ephemeral instances
    async fn heartbeat_loop(
        instances: Arc<RwLock<Vec<(String, String, Instance)>>>,
        rpc_client: Arc<RpcClient>,
        namespace: String,
        shutdown: Arc<Notify>,
    ) {
        let heartbeat_interval = Duration::from_secs(5);

        loop {
            tokio::select! {
                _ = shutdown.notified() => {
                    info!("Heartbeat loop shutdown");
                    break;
                }
                _ = tokio::time::sleep(heartbeat_interval) => {
                    let registered = instances.read().clone();

                    for (service_name, group_name, instance) in registered {
                        if !instance.ephemeral {
                            continue;
                        }

                        // Re-register to maintain heartbeat
                        let request = InstanceRequest::register(
                            &namespace,
                            &service_name,
                            &group_name,
                            instance,
                        );

                        if let Err(e) = rpc_client.request::<_, InstanceResponse>(&request).await {
                            warn!(
                                "Heartbeat failed for service={}, group={}: {}",
                                service_name, group_name, e
                            );
                        } else {
                            debug!(
                                "Heartbeat sent for service={}, group={}",
                                service_name, group_name
                            );
                        }
                    }
                }
            }
        }
    }
}

impl Drop for NamingService {
    fn drop(&mut self) {
        self.shutdown.notify_one();
    }
}

/// Simple random index generator
fn rand_index(max: usize) -> usize {
    use std::time::{SystemTime, UNIX_EPOCH};
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .subsec_nanos() as usize;
    nanos % max
}