quickdid / src / handle_resolver / memory.rs

//! In-memory caching handle resolver.
//!
//! This module provides a handle resolver that caches resolution results in memory
//! with a configurable TTL. This is useful for reducing DNS/HTTP lookups and
//! improving performance when Redis is not available.

use super::errors::HandleResolverError;
use super::traits::HandleResolver;
use crate::metrics::SharedMetricsPublisher;
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::RwLock;

/// Result of a handle resolution cached in memory.
#[derive(Clone, Debug)]
enum ResolveHandleResult {
    /// Handle was successfully resolved to a DID
    Found(u64, String),
    /// Handle resolution failed
    NotFound(u64, String),
}

/// In-memory caching wrapper for handle resolvers.
///
/// This resolver wraps another resolver and caches results in memory with
/// a configurable TTL. Both successful and failed resolutions are cached
/// to avoid repeated lookups.
///
/// # Example
///
/// ```no_run
/// use std::sync::Arc;
/// use quickdid::handle_resolver::{create_caching_resolver, create_base_resolver, HandleResolver};
/// use quickdid::metrics::NoOpMetricsPublisher;
///
/// # async fn example() {
/// # use atproto_identity::resolve::HickoryDnsResolver;
/// # use reqwest::Client;
/// # let dns_resolver = Arc::new(HickoryDnsResolver::create_resolver(&[]));
/// # let http_client = Client::new();
/// # let metrics = Arc::new(NoOpMetricsPublisher);
/// let base_resolver = create_base_resolver(dns_resolver, http_client, metrics.clone());
/// let caching_resolver = create_caching_resolver(
///     base_resolver,
///     300,  // 5 minute TTL
///     metrics
/// );
///
/// // First call hits the underlying resolver
/// let (did1, timestamp1) = caching_resolver.resolve("alice.bsky.social").await.unwrap();
///
/// // Second call returns cached result
/// let (did2, timestamp2) = caching_resolver.resolve("alice.bsky.social").await.unwrap();
/// # }
/// ```
pub(super) struct CachingHandleResolver {
    inner: Arc<dyn HandleResolver>,
    cache: Arc<RwLock<HashMap<String, ResolveHandleResult>>>,
    ttl_seconds: u64,
    metrics: SharedMetricsPublisher,
}

impl CachingHandleResolver {
    /// Create a new caching handle resolver.
    ///
    /// # Arguments
    ///
    /// * `inner` - The underlying resolver to use for actual resolution
    /// * `ttl_seconds` - How long to cache results in seconds
    /// * `metrics` - Metrics publisher for telemetry
    pub fn new(
        inner: Arc<dyn HandleResolver>,
        ttl_seconds: u64,
        metrics: SharedMetricsPublisher,
    ) -> Self {
        Self {
            inner,
            cache: Arc::new(RwLock::new(HashMap::new())),
            ttl_seconds,
            metrics,
        }
    }

    fn current_timestamp() -> u64 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs()
    }

    fn is_expired(&self, timestamp: u64) -> bool {
        let current = Self::current_timestamp();
        current > timestamp && (current - timestamp) > self.ttl_seconds
    }
}

#[async_trait]
impl HandleResolver for CachingHandleResolver {
    async fn resolve(&self, s: &str) -> Result<(String, u64), HandleResolverError> {
        let handle = s.to_string();

        // Check cache first
        {
            let cache = self.cache.read().await;
            if let Some(cached) = cache.get(&handle) {
                match cached {
                    ResolveHandleResult::Found(timestamp, did) => {
                        if !self.is_expired(*timestamp) {
                            tracing::debug!("Cache hit for handle {}: {}", handle, did);
                            self.metrics.incr("resolver.memory.cache_hit").await;
                            return Ok((did.clone(), *timestamp));
                        }
                        tracing::debug!("Cache entry expired for handle {}", handle);
                        self.metrics.incr("resolver.memory.cache_expired").await;
                    }
                    ResolveHandleResult::NotFound(timestamp, error) => {
                        if !self.is_expired(*timestamp) {
                            tracing::debug!(
                                "Cache hit (not found) for handle {}: {}",
                                handle,
                                error
                            );
                            self.metrics
                                .incr("resolver.memory.cache_hit_not_resolved")
                                .await;
                            return Err(HandleResolverError::HandleNotFoundCached(error.clone()));
                        }
                        tracing::debug!("Cache entry expired for handle {}", handle);
                        self.metrics.incr("resolver.memory.cache_expired").await;
                    }
                }
            }
        }

        // Not in cache or expired, resolve through inner resolver
        tracing::debug!("Cache miss for handle {}, resolving...", handle);
        self.metrics.incr("resolver.memory.cache_miss").await;
        let result = self.inner.resolve(s).await;

        // Store in cache
        {
            let mut cache = self.cache.write().await;
            match &result {
                Ok((did, timestamp)) => {
                    cache.insert(
                        handle.clone(),
                        ResolveHandleResult::Found(*timestamp, did.clone()),
                    );
                    self.metrics.incr("resolver.memory.cache_set").await;
                    tracing::debug!(
                        "Cached successful resolution for handle {}: {}",
                        handle,
                        did
                    );
                }
                Err(e) => {
                    let timestamp = Self::current_timestamp();
                    cache.insert(
                        handle.clone(),
                        ResolveHandleResult::NotFound(timestamp, e.to_string()),
                    );
                    self.metrics.incr("resolver.memory.cache_set_error").await;
                    tracing::debug!("Cached failed resolution for handle {}: {}", handle, e);
                }
            }

            // Track cache size
            let cache_size = cache.len() as u64;
            self.metrics
                .gauge("resolver.memory.cache_entries", cache_size)
                .await;
        }

        result
    }

    async fn set(&self, handle: &str, did: &str) -> Result<(), HandleResolverError> {
        // Normalize the handle to lowercase
        let handle = handle.to_lowercase();

        // Update the in-memory cache
        {
            let mut cache = self.cache.write().await;
            let timestamp = Self::current_timestamp();
            cache.insert(
                handle.clone(),
                ResolveHandleResult::Found(timestamp, did.to_string()),
            );
            self.metrics.incr("resolver.memory.set").await;
            tracing::debug!("Set handle {} -> DID {} in memory cache", handle, did);

            // Track cache size
            let cache_size = cache.len() as u64;
            self.metrics
                .gauge("resolver.memory.cache_entries", cache_size)
                .await;
        }

        // Chain to inner resolver
        self.inner.set(&handle, did).await
    }
}

/// Create a new in-memory caching handle resolver.
///
/// This factory function creates a resolver that caches resolution results
/// in memory with a configurable TTL.
///
/// # Arguments
///
/// * `inner` - The underlying resolver to use for actual resolution
/// * `ttl_seconds` - How long to cache results in seconds
/// * `metrics` - Metrics publisher for telemetry
///
/// # Example
///
/// ```no_run
/// use std::sync::Arc;
/// use quickdid::handle_resolver::{create_base_resolver, create_caching_resolver, HandleResolver};
/// use quickdid::metrics::NoOpMetricsPublisher;
///
/// # async fn example() {
/// # use atproto_identity::resolve::HickoryDnsResolver;
/// # use reqwest::Client;
/// # let dns_resolver = Arc::new(HickoryDnsResolver::create_resolver(&[]));
/// # let http_client = Client::new();
/// # let metrics = Arc::new(NoOpMetricsPublisher);
/// let base = create_base_resolver(
///     dns_resolver,
///     http_client,
///     metrics.clone(),
/// );
///
/// let resolver = create_caching_resolver(base, 300, metrics);  // 5 minute TTL
/// let did = resolver.resolve("alice.bsky.social").await.unwrap();
/// # }
/// ```
pub fn create_caching_resolver(
    inner: Arc<dyn HandleResolver>,
    ttl_seconds: u64,
    metrics: SharedMetricsPublisher,
) -> Arc<dyn HandleResolver> {
    Arc::new(CachingHandleResolver::new(inner, ttl_seconds, metrics))
}