Rustで一定時間データをキャッシュする仕組みを作った

目的

APIで大量にアクセスがある場合、あまり変わらない想定のマスターデータがAPIの呼び出しの都度データベースにアクセスしてほしくないです。

そこで一定時間データをキャッシュする仕組みを作りました。

データはハッシュマップに保存して、同じキーの場合、一定時間なら同じデータを返して、過ぎていた場合新しいデータを取ってきます。

ハッシュマップはonce_cellを使ってグローバルにおいておきます。

結果

ランダムな文字列を保持するModelがあります。これを1秒間キャッシュします。
以下の結果では最初の2回は同じ結果ですが、3秒後には違う結果になっています。

Some(Model { value: "3u0TUQ79MrSTY9VE" })
Some(Model { value: "3u0TUQ79MrSTY9VE" })

# 3秒経過

Some(Model { value: "hko6Tvjg1MT48A6u" })

コード

Cargo.toml

[package]
name = "holder_sample"
version = "0.1.0"
edition = "2021"

[dependencies]
once_cell = "1"
thiserror = "1"
chrono = { version = "0.4", features = ["serde", "clock"] }
tokio = { version = "1", features = ["full"] }
rand = "0.8.5"

main.rs

pub mod error;
pub mod holder;
pub mod model;

use crate::holder::Holder;
use crate::model::Model;
use once_cell::sync::Lazy;
use std::collections::HashMap;
use std::sync::Mutex;

// グローバルなデータをキャッシュするハッシュマップ
// キーはモデルを一意で表す文字列です。DBなら主キーなど。
static MODEL_HASH: Lazy<Mutex<HashMap<String, Holder<Model>>>> =
    Lazy::new(|| Mutex::new(HashMap::new()));

#[tokio::main]
async fn main() -> Result<(), crate::error::Error> {
    // キャッシュ保持期間は1秒
    let cache_expire_second_count = 1;

    // 最初のモデルを取得する。最初はキャッシュされていないので、クロージャーが呼ばれて生成される
    let model = Holder::get("12345", &MODEL_HASH, cache_expire_second_count, || async {
        Ok(Some(Model::new()))
    })
    .await?;
    println!("{:?}", model);

    // すぐに取得すると、キャッシュが利用されて上記と同じ結果になる
    // クロージャーは呼ばれない
    let model = Holder::get("12345", &MODEL_HASH, cache_expire_second_count, || async {
        Ok(Some(Model::new()))
    })
    .await?;
    println!("{:?}", model);

    // しばらく待つ
    tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;

    // キャッシュ保持期間が過ぎるのでクロージャーが呼ばれて新しいデータができる
    let model = Holder::get("12345", &MODEL_HASH, cache_expire_second_count, || async {
        Ok(Some(Model::new()))
    })
    .await?;
    println!("{:?}", model);

    Ok(())
}

error.rs

use thiserror::Error;

#[derive(Error, Debug)]
pub enum Error {
    #[error("Invalid {0}")]
    Invalid(String),
}

model.rs

use rand::distributions::{Alphanumeric, DistString};

#[derive(Debug, Clone)]
pub struct Model {
    pub value: String,
}

impl Model {
    // ランダムな文字列で生成する
    pub fn new() -> Self {
        let mut rng = rand::thread_rng();
        Self {
            value: Alphanumeric.sample_string(&mut rng, 16),
        }
    }
}

holder.rs

use crate::error::Error;
pub use chrono::{prelude::*, Duration};
use std::{collections::HashMap, future::Future, sync::Mutex};

pub struct Holder<T> {
    // 保持するデータ
    target: T,
    // キャッシュが無効になる時間
    expire_at: DateTime<Utc>,
}

impl<T> Holder<T>
where
    T: Clone,
{
    pub fn new(target: T, expire_at: DateTime<Utc>) -> Self {
        Self { target, expire_at }
    }

    fn get_target(&self, now: &DateTime<Utc>) -> Option<T> {
        if self.expire_at > *now {
            Some(self.target.clone())
        } else {
            None
        }
    }

    fn inner_get(
        key: &str,
        now: &DateTime<Utc>,
        mutex: &Mutex<HashMap<String, Holder<T>>>,
    ) -> Result<Option<T>, Error> {
        let hash = mutex
            .lock()
            .map_err(|_| Error::Invalid("lock failure".to_owned()))?;
        let holder = hash.get(key);
        Ok(if let Some(holder) = holder.as_ref() {
            holder.get_target(now)
        } else {
            None
        })
    }

    fn inner_set(
        key: &str,
        target: T,
        mutex: &Mutex<HashMap<String, Holder<T>>>,
        expire_at: DateTime<Utc>,
    ) -> Result<(), Error> {
        let mut hash = mutex
            .lock()
            .map_err(|_| Error::Invalid("lock failure".to_owned()))?;
        hash.insert(key.to_owned(), Holder::new(target, expire_at));
        Ok(())
    }

    pub async fn get<Fut>(
        key: &str,
        mutex: &Mutex<HashMap<String, Holder<T>>>,
        cache_expire_second_count: u64,
        f: impl FnOnce() -> Fut,
    ) -> Result<Option<T>, Error>
    where
        Fut: Future<Output = Result<Option<T>, Error>>,
    {
        let now = Utc::now();
        {
            if let Some(target) = Self::inner_get(key, &now, mutex)? {
                return Ok(Some(target));
            }
        }
        let target = f().await?;
        if let Some(target) = target.as_ref() {
            Self::inner_set(
                key,
                target.clone(),
                mutex,
                now + Duration::seconds(cache_expire_second_count as i64),
            )?;
        }
        Ok(target)
    }
}