photoncloud-monorepo/mtls-agent/src/discovery.rs

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::Result;
use chainfire_client::Client;
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::{info, warn};

const PHOTON_PREFIX: &str = "photoncloud";
const CACHE_TTL: Duration = Duration::from_secs(30);
const POLICY_CACHE_TTL: Duration = Duration::from_secs(30);

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServiceInstance {
    pub instance_id: String,
    pub service: String,
    pub node_id: String,
    pub ip: String,
    pub port: u16,
    #[serde(default)]
    pub mesh_port: Option<u16>,
    #[serde(default)]
    pub version: Option<String>,
    #[serde(default)]
    pub state: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MtlsPolicy {
    pub policy_id: String,
    #[serde(default)]
    pub environment: Option<String>,
    pub source_service: String,
    pub target_service: String,
    #[serde(default)]
    pub mtls_required: Option<bool>,
    #[serde(default)]
    pub mode: Option<String>,
}

struct CachedInstances {
    instances: Vec<ServiceInstance>,
    updated_at: Instant,
}

struct CachedPolicy {
    policy: MtlsPolicy,
    updated_at: Instant,
}

pub struct ServiceDiscovery {
    chainfire_endpoint: String,
    cluster_id: String,
    cache: Arc<RwLock<HashMap<String, CachedInstances>>>,
    policy_cache: Arc<RwLock<HashMap<String, CachedPolicy>>>,
}

impl ServiceDiscovery {
    pub fn new(chainfire_endpoint: String, cluster_id: String) -> Self {
        Self {
            chainfire_endpoint,
            cluster_id,
            cache: Arc::new(RwLock::new(HashMap::new())),
            policy_cache: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn resolve_service(&self, service_name: &str) -> Result<Vec<ServiceInstance>> {
        // キャッシュをチェック
        {
            let cache = self.cache.read().await;
            if let Some(cached) = cache.get(service_name) {
                if cached.updated_at.elapsed() < CACHE_TTL {
                    return Ok(cached.instances.clone());
                }
            }
        }

        // Chainfireから取得
        let instances = self.fetch_instances_from_chainfire(service_name).await?;

        // キャッシュを更新
        {
            let mut cache = self.cache.write().await;
            cache.insert(
                service_name.to_string(),
                CachedInstances {
                    instances: instances.clone(),
                    updated_at: Instant::now(),
                },
            );
        }

        Ok(instances)
    }

    async fn fetch_instances_from_chainfire(&self, service_name: &str) -> Result<Vec<ServiceInstance>> {
        let mut client = Client::connect(self.chainfire_endpoint.clone()).await?;
        let prefix = format!(
            "{}instances/{}/",
            cluster_prefix(&self.cluster_id),
            service_name
        );
        let prefix_bytes = prefix.as_bytes();

        let (kvs, _) = client.scan_prefix(prefix_bytes, 0).await?;
        let mut instances = Vec::new();

        for (_, value, _) in kvs {
            match serde_json::from_slice::<ServiceInstance>(&value) {
                Ok(inst) => {
                    // 状態が "healthy" または未設定のもののみ返す
                    if inst.state.as_deref().unwrap_or("healthy") == "healthy" {
                        instances.push(inst);
                    }
                }
                Err(e) => {
                    warn!(error = %e, "failed to parse ServiceInstance from Chainfire");
                }
            }
        }

        info!(
            service = %service_name,
            count = instances.len(),
            "resolved service instances from Chainfire"
        );

        Ok(instances)
    }

    pub async fn get_mtls_policy(
        &self,
        source_service: &str,
        target_service: &str,
    ) -> Result<Option<MtlsPolicy>> {
        let policy_key = format!(
            "{}-{}",
            source_service, target_service
        );

        // キャッシュをチェック
        {
            let cache = self.policy_cache.read().await;
            if let Some(cached) = cache.get(&policy_key) {
                if cached.updated_at.elapsed() < POLICY_CACHE_TTL {
                    return Ok(Some(cached.policy.clone()));
                }
            }
        }

        // Chainfireから取得
        let mut client = Client::connect(self.chainfire_endpoint.clone()).await?;
        let prefix = format!(
            "{}mtls/policies/",
            cluster_prefix(&self.cluster_id)
        );
        let prefix_bytes = prefix.as_bytes();

        let (kvs, _) = client.scan_prefix(prefix_bytes, 0).await?;

        for (_, value, _) in kvs {
            match serde_json::from_slice::<MtlsPolicy>(&value) {
                Ok(policy) => {
                    if policy.source_service == source_service && policy.target_service == target_service {
                        // キャッシュに保存
                        let mut cache = self.policy_cache.write().await;
                        cache.insert(
                            policy_key.clone(),
                            CachedPolicy {
                                policy: policy.clone(),
                                updated_at: Instant::now(),
                            },
                        );
                        return Ok(Some(policy));
                    }
                }
                Err(e) => {
                    warn!(error = %e, "failed to parse MtlsPolicy from Chainfire");
                }
            }
        }

        Ok(None)
    }

    pub async fn start_background_refresh(&self) {
        let endpoint = self.chainfire_endpoint.clone();
        let cluster_id = self.cluster_id.clone();
        let cache = Arc::clone(&self.cache);
        let policy_cache = Arc::clone(&self.policy_cache);

        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(60));
            loop {
                interval.tick().await;

                // 全サービスのインスタンスをリフレッシュ
                if let Ok(mut client) = Client::connect(endpoint.clone()).await {
                    let prefix = format!("{}instances/", cluster_prefix(&cluster_id));
                    if let Ok((kvs, _)) = client.scan_prefix(prefix.as_bytes(), 0).await {
                        let mut service_map: HashMap<String, Vec<ServiceInstance>> = HashMap::new();
                        for (_key, value, _) in kvs {
                            if let Ok(inst) = serde_json::from_slice::<ServiceInstance>(&value) {
                                if inst.state.as_deref().unwrap_or("healthy") == "healthy" {
                                    service_map
                                        .entry(inst.service.clone())
                                        .or_insert_with(Vec::new)
                                        .push(inst);
                                }
                            }
                        }
                        let mut cache_guard = cache.write().await;
                        for (service, instances) in service_map {
                            cache_guard.insert(
                                service,
                                CachedInstances {
                                    instances,
                                    updated_at: Instant::now(),
                                },
                            );
                        }
                    }
                }

                // ポリシーキャッシュはTTLベースでクリア
                {
                    let mut policy_guard = policy_cache.write().await;
                    policy_guard.retain(|_, cached| cached.updated_at.elapsed() < POLICY_CACHE_TTL);
                }
            }
        });
    }
}

fn cluster_prefix(cluster_id: &str) -> String {
    format!("{}/clusters/{}/", PHOTON_PREFIX, cluster_id)
}