有効期限付きの lru_cache

Posted by rhoboro on 2021-06-01

SQSにメッセージを送信する処理を書いていたのですが、結構な頻度で呼び出されるため boto3.client のインスタンスをキャッシュすることにしました。

Pythonではキャッシュといえばfunctools.lru_cacheを使うと楽なので、これをラップして有効期限を指定できるようにします。

有効期限付きの lru_cache

作ったコードはこちら。

from functools import lru_cache, wraps
from datetime import datetime, timedelta
from threading import Lock

def cache(seconds: int, max_size: int = 128, typed: bool = False):
    def wrapper(f):
        func = lru_cache(maxsize=max_size, typed=typed)(f)
        func.ttl = seconds
        func.expire = datetime.utcnow() + timedelta(seconds=func.ttl)

        @wraps(f)
        def inner(*args, **kwargs):
            with Lock():
                if datetime.utcnow() > func.expire:
                    func.cache_clear()
                    func.expire = datetime.utcnow() + timedelta(seconds=func.ttl)
                return func(*args, **kwargs)

        inner.clear_cache = func.cache_clear
        inner.cache_info = func.cache_info
        return inner

    return wrapper

念の為スレッドセーフにしています。

これを cache.py に保存した場合、こんな感じで利用できます。

$ python3 -q
>>> from cache import cache
>>> from datetime import datetime
>>>
>>> @cache(seconds=20)
... def add(x, y):
...   print(f"Expired!!!: {datetime.utcnow().isoformat()}")
...   return x + y
...
>>> add(1, 2)
Expired!!!: 2021-06-01T08:47:39.449230
3
>>> add(1, 2)
3
>>> add(1, 2)
3
>>> add(1, 2)
3
>>> add(1, 2)
3
>>> add(1, 2)
3
>>> add(1, 2)
Expired!!!: 2021-06-01T08:47:59.368536
3
>>> add(1, 2)
3

うまく動いているようです。

boto3.client での利用

実際のコードではこのようにしています。 ここでは boto3 を使っていますが、 google-cloud-bigquery などの Client でも全く同様にかけます。

import uuid

import boto3
from ... import cache


@cache(seconds=300)
def get_client(region_name):
    return boto3.client("sqs", region_name=region_name)

def send_message(message, queue_url, region_name):
    client = get_client(region_name)
    return client.send_message(
        QueueUrl=queue_url,
        MessageGroupId="sample",
        MessageDeduplicationId=str(uuid.uuid4()),
        MessageBody=message,
    )

本来は有効期限が切れてエラーになった時にリフレッシュできると良いのですが、サクッとできるAPIはなさそうだったので一旦この方針でいくことにしました。

tags: python