grafos_registry/

region.rs

//! Registry storage region backed by a FabricHashMap.

extern crate alloc;
use alloc::string::String;
use alloc::vec::Vec;

use grafos_collections::map::FabricHashMap;
use grafos_std::error::Result;
use grafos_std::host;
use serde::{Deserialize, Serialize};

use crate::filter::RegistryFilter;
use crate::registration::{HealthStatus, ServiceRegistration};

/// A slot in the registry map value — wraps a single registration.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct RegistrationSlot {
    pub registration: ServiceRegistration,
}

/// Fabric-backed registry storage.
///
/// Internally stores a `FabricHashMap<String, Vec<RegistrationSlot>>` where
/// the key is the service name and the value is a list of instance
/// registrations. A monotonic `version` counter is bumped on every mutation
/// to support [`RegistryWatcher`](crate::RegistryWatcher) change detection.
pub struct RegistryRegion {
    map: FabricHashMap<String, Vec<RegistrationSlot>>,
    version: u64,
    default_ttl_secs: u32,
}

impl RegistryRegion {
    /// Create a new registry region with the given capacity and strides.
    pub fn new(
        capacity: usize,
        default_ttl_secs: u32,
        name_stride: usize,
        slot_stride: usize,
    ) -> Result<Self> {
        let map = FabricHashMap::with_capacity(capacity, name_stride, slot_stride)?;
        Ok(RegistryRegion {
            map,
            version: 0,
            default_ttl_secs,
        })
    }

    /// Return the current version counter.
    pub fn version(&self) -> u64 {
        self.version
    }

    /// Register a service instance. If a registration with the same
    /// `instance_id` already exists under the same name, it is replaced.
    pub fn register(&mut self, mut reg: ServiceRegistration) -> Result<()> {
        let now = host::unix_time_secs();
        reg.lease_expires_at = now + self.default_ttl_secs as u64;

        let name = reg.name.clone();
        let mut slots = self.map.get(&name)?.unwrap_or_default();

        // Replace existing registration with same instance_id
        slots.retain(|s| s.registration.instance_id != reg.instance_id);
        slots.push(RegistrationSlot { registration: reg });

        self.map.insert(&name, &slots)?;
        self.version += 1;

        #[cfg(feature = "observe")]
        {
            let ver = &slots.last().unwrap().registration.version;
            crate::observe_hooks::on_service_registered(&name, ver);
        }

        Ok(())
    }

    /// Remove a specific instance from the registry.
    pub fn deregister(&mut self, name: &str, instance_id: u128) -> Result<bool> {
        let key = String::from(name);
        let mut slots = self.map.get(&key)?.unwrap_or_default();
        let before = slots.len();
        slots.retain(|s| s.registration.instance_id != instance_id);
        let removed = slots.len() < before;

        if slots.is_empty() {
            self.map.remove(&key)?;
        } else {
            self.map.insert(&key, &slots)?;
        }

        if removed {
            self.version += 1;
            #[cfg(feature = "observe")]
            crate::observe_hooks::on_service_deregistered(name);
        }
        Ok(removed)
    }

    /// Update the health status of a specific instance.
    pub fn set_health(
        &mut self,
        name: &str,
        instance_id: u128,
        health: HealthStatus,
    ) -> Result<bool> {
        let key = String::from(name);
        let mut slots = match self.map.get(&key)? {
            Some(s) => s,
            None => return Ok(false),
        };

        let mut found = false;
        for slot in &mut slots {
            if slot.registration.instance_id == instance_id {
                slot.registration.health = health.clone();
                found = true;
                break;
            }
        }

        if found {
            self.map.insert(&key, &slots)?;
            self.version += 1;
            #[cfg(feature = "observe")]
            {
                let old_str = "unknown";
                let new_str = match &health {
                    HealthStatus::Healthy => "healthy",
                    HealthStatus::Degraded { .. } => "degraded",
                    HealthStatus::Draining => "draining",
                };
                crate::observe_hooks::on_health_change(name, old_str, new_str);
            }
        }
        Ok(found)
    }

    /// Set a specific instance to [`HealthStatus::Draining`].
    pub fn set_draining(&mut self, name: &str, instance_id: u128) -> Result<bool> {
        self.set_health(name, instance_id, HealthStatus::Draining)
    }

    /// Look up all registrations for a service name.
    pub fn lookup(&self, name: &str) -> Result<Vec<ServiceRegistration>> {
        let key = String::from(name);
        match self.map.get(&key)? {
            Some(slots) => Ok(slots.into_iter().map(|s| s.registration).collect()),
            None => Ok(Vec::new()),
        }
    }

    /// Look up a single registration for a service name (returns the first match).
    pub fn lookup_one(&self, name: &str) -> Result<Option<ServiceRegistration>> {
        let results = self.lookup(name)?;
        Ok(results.into_iter().next())
    }

    /// Look up registrations matching a filter.
    pub fn lookup_filtered(
        &self,
        name: &str,
        filter: &RegistryFilter,
    ) -> Result<Vec<ServiceRegistration>> {
        let all = self.lookup(name)?;
        Ok(all.into_iter().filter(|r| filter.matches(r)).collect())
    }

    /// List all distinct service names in the registry.
    pub fn list_services(&self) -> Result<Vec<String>> {
        let mut names = Vec::new();
        for item in self.map.iter() {
            let (name, slots) = item?;
            if !slots.is_empty() {
                names.push(name);
            }
        }
        Ok(names)
    }

    /// Prune expired registrations. Returns the number of instances removed.
    pub fn tick(&mut self) -> Result<usize> {
        let now = host::unix_time_secs();
        let mut total_pruned = 0;

        // Collect all keys first to avoid borrow conflicts
        let mut keys = Vec::new();
        for item in self.map.iter() {
            let (name, _) = item?;
            keys.push(name);
        }

        for key in keys {
            let mut slots = match self.map.get(&key)? {
                Some(s) => s,
                None => continue,
            };
            let before = slots.len();
            slots.retain(|s| s.registration.lease_expires_at > now);
            let pruned = before - slots.len();

            if pruned > 0 {
                if slots.is_empty() {
                    self.map.remove(&key)?;
                } else {
                    self.map.insert(&key, &slots)?;
                }
                total_pruned += pruned;
            }
        }

        if total_pruned > 0 {
            self.version += 1;
        }
        Ok(total_pruned)
    }
}