gitks/server/cache.rs

use std::num::NonZeroUsize;
use std::sync::{Mutex, OnceLock};

use clru::CLruCache;
use prost::Message;

use crate::pb::{ObjectSelector, object_selector};

const GLOBAL_CACHE_MAX: usize = 65_545;

type Cache = CLruCache<Vec<u8>, Vec<u8>>;

static GLOBAL_CACHE: OnceLock<Mutex<Cache>> = OnceLock::new();

fn cache() -> &'static Mutex<Cache> {
    GLOBAL_CACHE.get_or_init(|| {
        let capacity =
            NonZeroUsize::new(GLOBAL_CACHE_MAX).expect("cache capacity must be non-zero");
        Mutex::new(CLruCache::new(capacity))
    })
}

fn cache_key<Req>(namespace: &str, request: &Req) -> Vec<u8>
where
    Req: Message,
{
    let mut key = Vec::with_capacity(namespace.len() + 1 + request.encoded_len());
    key.extend_from_slice(namespace.as_bytes());
    key.push(0);
    request
        .encode(&mut key)
        .expect("encoding a prost message into Vec cannot fail");
    key
}

pub(crate) fn cached_response<Req, Res, E, F>(
    namespace: &'static str,
    request: &Req,
    build: F,
) -> Result<Res, E>
where
    Req: Message,
    Res: Message + Default,
    F: FnOnce() -> Result<Res, E>,
{
    let key = cache_key(namespace, request);

    if let Some(bytes) = cache()
        .lock()
        .unwrap_or_else(|e| e.into_inner())
        .get(&key)
        .cloned()
        && let Ok(response) = Res::decode(bytes.as_slice())
    {
        tracing::debug!(
            namespace = %namespace,
            key_len = key.len(),
            "cache hit"
        );
        return Ok(response);
    }

    tracing::debug!(
        namespace = %namespace,
        key_len = key.len(),
        "cache miss, building response"
    );
    let response = build()?;
    let mut bytes = Vec::with_capacity(response.encoded_len());
    response
        .encode(&mut bytes)
        .expect("encoding a prost message into Vec cannot fail");
    cache()
        .lock()
        .unwrap_or_else(|e| e.into_inner())
        .put(key, bytes);
    Ok(response)
}

pub(crate) fn cached_vec_response<Req, Item, E, F>(
    namespace: &'static str,
    request: &Req,
    build: F,
) -> Result<Vec<Item>, E>
where
    Req: Message,
    Item: Message + Default,
    F: FnOnce() -> Result<Vec<Item>, E>,
{
    let key = cache_key(namespace, request);

    if let Some(bytes) = cache()
        .lock()
        .unwrap_or_else(|e| e.into_inner())
        .get(&key)
        .cloned()
    {
        let mut remaining = bytes.as_slice();
        let mut items = Vec::new();
        let mut valid = true;
        while !remaining.is_empty() {
            match Item::decode_length_delimited(&mut remaining) {
                Ok(item) => items.push(item),
                Err(_) => {
                    valid = false;
                    break;
                }
            }
        }
        if valid {
            tracing::debug!(
                namespace = %namespace,
                key_len = key.len(),
                item_count = items.len(),
                "vec cache hit"
            );
            return Ok(items);
        }
        tracing::warn!(
            namespace = %namespace,
            "vec cache decode failed, rebuilding"
        );
    }

    tracing::debug!(
        namespace = %namespace,
        key_len = key.len(),
        "vec cache miss, building response"
    );
    let response = build()?;
    let mut bytes = Vec::new();
    for item in &response {
        item.encode_length_delimited(&mut bytes)
            .expect("encoding a prost message into Vec cannot fail");
    }
    cache()
        .lock()
        .unwrap_or_else(|e| e.into_inner())
        .put(key, bytes);
    Ok(response)
}

pub(crate) fn selector_is_oid(selector: &Option<ObjectSelector>) -> bool {
    matches!(
        selector.as_ref().and_then(|s| s.selector.as_ref()),
        Some(object_selector::Selector::Oid(_))
    )
}

pub(crate) fn selectors_are_oid(
    left: &Option<ObjectSelector>,
    right: &Option<ObjectSelector>,
) -> bool {
    selector_is_oid(left) && selector_is_oid(right)
}