RustでPostgreSQLのデータを個別の破棄時間でキャッシュする。

目的

RustでPostgreSQLのデータをキャッシュするの記事では、DBのデータをキャッシュして、一定時間が過ぎたら全て破棄する仕組みでした。

今回の記事では、それぞれの値が取得してからの破棄時間を見て取得し直す仕組みにしています。

コード

holder

今回は破棄時間をHashMapの値のタプルの中に保存しています。
これを使ってキーごとに破棄時間を確認します。

pub struct HolderMapEachExpire<K, V> {
    map: HashMap<K, (V, DateTime<Utc>)>,
    expire_interval: Duration,
    pg_pool: deadpool_postgres::Pool,
}

impl<K, V> HolderMapEachExpire<K, V>
where
    K: PartialEq + Eq + Hash + Clone,
    V: Clone,
{
    pub fn new(pg_pool: deadpool_postgres::Pool, expire_interval: Duration) -> Self {
        Self {
            map: HashMap::new(),
            expire_interval,
            pg_pool,
        }
    }

    pub async fn get<FutOne>(
        &mut self,
        key: &K,
        now: Option<DateTime<Utc>>,
        f: impl FnOnce(deadpool_postgres::Client, K) -> FutOne,
    ) -> Result<Option<V>, Error>
    where
        FutOne: Future<Output = Result<Option<V>, Error>>,
    {
        match self.map.get(key) {
            Some((value, expire_at)) if get_now(now) < *expire_at => {
                return Ok(Some(value.clone()));
            }
            _ => {}
        }
        let pg_client = self.pg_pool.get().await?;
        let Some(value) = f(pg_client, key.clone()).await? else {
            return Ok(None);
        };
        self.map.insert(
            key.clone(),
            (value.clone(), expire_at(now, self.expire_interval)),
        );
        Ok(Some(value))
    }
}

fn get_now(now: Option<DateTime<Utc>>) -> DateTime<Utc> {
    now.unwrap_or(Utc::now())
}

fn expire_at(now: Option<DateTime<Utc>>, interval: Duration) -> DateTime<Utc> {
    get_now(now) + interval
}

main

テストでは2回連続取得して、最初だけDBが呼び出されていることを確認します。
4秒まってもう一度取得した時にDBが呼び出されているか確認します。

async fn get_account_each_expire(uuid: &Uuid) -> anyhow::Result<Option<Account>> {
    let mut holder = HOLDER_EACHEXPIRE.get().unwrap().lock().await;
    let account = holder
        .get(uuid, None, |pg_client, uuid| async move {
            println!("each expire one");
            Account::get_one(&pg_client, &uuid).await
        })
        .await
        .unwrap();
    Ok(account)
}


#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let pg_url =
        std::env::var("PG_URL").unwrap_or("postgres://user:pass@localhost:5432/web".to_owned());
    let pg_pool = get_postgres_pool(&pg_url)?;

    // スレッドが始まる前に初期化する
    let _ = HOLDER_EACHEXPIRE.get_or_init(|| {
        Mutex::new(HolderMapEachExpire::new(
            pg_pool.clone(),
            Duration::from_secs(3),
        ))
    });

    let thread = tokio::spawn(async move {
        let account = get_account_each_expire(
            &Uuid::from_str("00000000-0000-0000-0000-000000000001").unwrap(),
        )
        .await
        .unwrap();
        println!("{:?}", account);
        let account = get_account_each_expire(
            &Uuid::from_str("00000000-0000-0000-0000-000000000001").unwrap(),
        )
        .await
        .unwrap();
        println!("{:?}", account);
        sleep(Duration::from_secs(4)).await;
        let account = get_account_each_expire(
            &Uuid::from_str("00000000-0000-0000-0000-000000000001").unwrap(),
        )
        .await
        .unwrap();
        println!("{:?}", account);
    });

    thread.await?;

    Ok(())
}

結果

each expire one
Some(Account { uuid: 00000000-0000-0000-0000-000000000001, content: "1" })
Some(Account { uuid: 00000000-0000-0000-0000-000000000001, content: "1" })
each expire one
Some(Account { uuid: 00000000-0000-0000-0000-000000000001, content: "1" })

まとめ

今回のコードもcrate resident-utils
でリリースしています。
サンプルコードはここにあります。

前回と今回とで全体の破棄時間と個別の破棄時間を考慮したキャッシュができるようになりました。