kftray-portforward 0.26.7

KFtray library with port forwarding logic
Documentation
use std::sync::Arc;
use std::time::{
    Duration,
    Instant,
};

use dashmap::DashMap;
use kube::Client;
use once_cell::sync::Lazy;
use tokio::sync::Mutex;

use crate::kube::client::create_client_with_specific_context;

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct ServiceClientKey {
    pub context_name: Option<String>,
    pub kubeconfig_path: Option<String>,
    pub config_id: i64,
}

impl ServiceClientKey {
    pub fn new(
        context_name: Option<String>, kubeconfig_path: Option<String>, config_id: i64,
    ) -> Self {
        Self {
            context_name,
            kubeconfig_path,
            config_id,
        }
    }
}

struct CachedClient {
    client: Arc<Client>,
    created_at: Instant,
}

impl CachedClient {
    fn new(client: Client) -> Self {
        Self {
            client: Arc::new(client),
            created_at: Instant::now(),
        }
    }

    fn is_expired(&self, ttl: Duration) -> bool {
        self.created_at.elapsed() > ttl
    }
}

pub struct SharedClientManager {
    clients: DashMap<ServiceClientKey, CachedClient>,
    client_ttl: Duration,
    cache_lock: Mutex<()>,
}

impl SharedClientManager {
    pub fn new() -> Self {
        Self {
            clients: DashMap::new(),
            client_ttl: Duration::from_secs(3600),
            cache_lock: Mutex::new(()),
        }
    }

    pub async fn get_client(&self, key: ServiceClientKey) -> anyhow::Result<Arc<Client>> {
        let _lock = self.cache_lock.lock().await;

        if let Some(cached) = self.clients.get(&key) {
            if !cached.is_expired(self.client_ttl) {
                return Ok(cached.client.clone());
            }
            self.clients.remove(&key);
        }

        let (client_opt, _, _) = create_client_with_specific_context(
            key.kubeconfig_path.clone(),
            key.context_name.as_deref(),
        )
        .await?;

        let client = client_opt.ok_or_else(|| {
            anyhow::anyhow!(
                "Failed to create client for context: {:?}",
                key.context_name
            )
        })?;

        let cached_client = CachedClient::new(client);
        let client_arc = cached_client.client.clone();
        self.clients.insert(key, cached_client);
        Ok(client_arc)
    }

    pub fn invalidate_client(&self, key: &ServiceClientKey) {
        self.clients.remove(key);
    }

    pub fn cleanup_expired(&self) {
        self.clients
            .retain(|_, cached| !cached.is_expired(self.client_ttl));
    }
}

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

pub static SHARED_CLIENT_MANAGER: Lazy<SharedClientManager> = Lazy::new(SharedClientManager::new);

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_service_client_key() {
        let key1 = ServiceClientKey::new(
            Some("context1".to_string()),
            Some("/path/to/config".to_string()),
            1,
        );

        let key2 = ServiceClientKey::new(
            Some("context1".to_string()),
            Some("/path/to/config".to_string()),
            1,
        );

        assert_eq!(key1, key2);
    }
}