kapiti 0.0.3

The Kapiti DNS Server
Documentation
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result};
use async_trait::async_trait;
use tracing::{debug, info, warn};

use crate::cache::DnsCache;
use crate::codec::message::{self, RequestInfo};
use crate::specs::message::Message;

/// Client for fetching and storing DNS lookup results in a local Retainer cache
pub struct Cache {
    cache: Arc<retainer::cache::Cache<String, Message>>,
    _monitor: smol::Task<()>,
    max_records: usize,
}

impl Cache {
    pub fn new(max_records: usize) -> Cache {
        info!("Using in-memory cache with max_records={}", max_records);
        let cache = Arc::new(retainer::cache::Cache::new());
        let cache_clone = cache.clone();
        // Set up background job to clean expired entries
        let monitor = smol::spawn(async move {
            // Every 3s, check 4 entries for expiration.
            // If >25% were expired, then repeat the check for another 4 entries
            cache_clone.monitor(4, 0.25, Duration::from_secs(3)).await
        });
        Cache {
            cache,
            _monitor: monitor,
            max_records,
        }
    }

    fn key(self: &Cache, request_info: &RequestInfo) -> String {
        super::cache_key(request_info)
    }
}

#[async_trait]
impl DnsCache for Cache {
    /// Queries Redis for a cached response.
    async fn fetch(&mut self, request_info: RequestInfo) -> Result<Option<Message>> {
        let cache_key = self.key(&request_info);
        match self.cache.get(&cache_key).await {
            Some(guard) => {
                // remaining() calculates Instant::now() internally. If that someday
                // became a problem we could call instant() to get the underlying expire time.
                let cache_ttl = guard.expiration().remaining().with_context(|| {
                    format!("Cache entry lacks expiration for cache_key='{}'", cache_key)
                })?;
                let mut updated_response = (*guard).clone();
                debug!(
                    "Cached response for cache_key='{}': (ttl={:?}) {}",
                    cache_key, cache_ttl, updated_response
                );
                if let Err(e) = message::update_cached_response(
                    &mut updated_response,
                    &request_info,
                    cache_ttl.as_secs() as u32,
                ) {
                    warn!(
                        "Ignoring cache response at cache_key='{}': {}",
                        cache_key, e
                    );
                    // Remove the bad entry before returning
                    self.cache.remove(&cache_key).await;
                    Ok(None)
                } else {
                    Ok(Some(updated_response))
                }
            }
            None => {
                debug!(
                    "Cache didn't have cached {:?} result for {}",
                    request_info.resource_type, request_info.name
                );
                Ok(None)
            }
        }
    }

    /// Stores an upstream server response Message to the cache.
    async fn store(&mut self, request_info: RequestInfo, response: Message) -> Result<()> {
        // Before storing, check if the cache has reached the max record count, and clear it if so.
        // This avoids unbounded storage and favors new retrievals over old cached values with long TTLs.
        if self.max_records > 0 {
            let cache_len = self.cache.len().await;
            if cache_len >= self.max_records {
                warn!(
                    "Clearing local cache: size={} max_records={}",
                    cache_len, self.max_records
                );
                self.cache.clear().await;
            }
        }

        match message::get_min_ttl_secs(&response) {
            // This can happen with e.g. a SERVFAIL response
            None => debug!(
                "Skipping storage of {:?} response for {} with missing resources",
                request_info.resource_type, request_info.name
            ),
            // Not much point in storing something that expires right away
            // (Might not happen in practice?)
            Some(0) => debug!(
                "Skipping storage of {:?} result for {} with TTL=0s",
                request_info.resource_type, request_info.name
            ),
            Some(response_min_ttl) => {
                let cache_key = self.key(&request_info);
                debug!(
                    "Stored response for {:?} {} request to cache_key='{}' with TTL={}s",
                    request_info.resource_type, request_info.name, cache_key, response_min_ttl
                );
                self.cache
                    .insert(
                        cache_key,
                        response,
                        Duration::from_secs(response_min_ttl as u64),
                    )
                    .await;
            }
        }
        Ok(())
    }
}