NDOJ/judge/caching.py

from inspect import signature
from django.core.cache import cache, caches
from django.db.models.query import QuerySet
from django.core.handlers.wsgi import WSGIRequest

import hashlib

from judge.logging import log_debug

MAX_NUM_CHAR = 50
NONE_RESULT = "__None__"


def arg_to_str(arg):
    if hasattr(arg, "id"):
        return str(arg.id)
    if isinstance(arg, list) or isinstance(arg, QuerySet):
        return hashlib.sha1(str(list(arg)).encode()).hexdigest()[:MAX_NUM_CHAR]
    if len(str(arg)) > MAX_NUM_CHAR:
        return str(arg)[:MAX_NUM_CHAR]
    return str(arg)


def filter_args(args_list):
    return [x for x in args_list if not isinstance(x, WSGIRequest)]


l0_cache = caches["l0"] if "l0" in caches else None


def cache_wrapper(prefix, timeout=None, expected_type=None):
    def get_key(func, *args, **kwargs):
        args_list = list(args)
        signature_args = list(signature(func).parameters.keys())
        args_list += [kwargs.get(k) for k in signature_args[len(args) :]]
        args_list = filter_args(args_list)
        args_list = [arg_to_str(i) for i in args_list]
        key = prefix + ":" + ":".join(args_list)
        key = key.replace(" ", "_")
        return key

    def _get(key):
        if not l0_cache:
            return cache.get(key)
        result = l0_cache.get(key)
        if result is None:
            result = cache.get(key)
        return result

    def _set_l0(key, value):
        if l0_cache:
            l0_cache.set(key, value, 30)

    def _set(key, value, timeout):
        _set_l0(key, value)
        cache.set(key, value, timeout)

    def decorator(func):
        def _validate_type(cache_key, result):
            if expected_type and not isinstance(result, expected_type):
                data = {
                    "function": f"{func.__module__}.{func.__qualname__}",
                    "result": str(result)[:30],
                    "expected_type": expected_type,
                    "type": type(result),
                    "key": cache_key,
                }
                log_debug("invalid_key", data)
                return False
            return True

        def wrapper(*args, **kwargs):
            cache_key = get_key(func, *args, **kwargs)
            result = _get(cache_key)
            if result is not None and _validate_type(cache_key, result):
                _set_l0(cache_key, result)
                if type(result) == str and result == NONE_RESULT:
                    result = None
                return result
            result = func(*args, **kwargs)
            if result is None:
                result = NONE_RESULT
            _set(cache_key, result, timeout)
            return result

        def dirty(*args, **kwargs):
            cache_key = get_key(func, *args, **kwargs)
            cache.delete(cache_key)
            if l0_cache:
                l0_cache.delete(cache_key)

        def prefetch_multi(args_list):
            keys = []
            for args in args_list:
                keys.append(get_key(func, *args))
            results = cache.get_many(keys)
            for key, result in results.items():
                if result is not None:
                    _set_l0(key, result)

        def dirty_multi(args_list):
            keys = []
            for args in args_list:
                keys.append(get_key(func, *args))
            cache.delete_many(keys)
            if l0_cache:
                l0_cache.delete_many(keys)

        wrapper.dirty = dirty
        wrapper.prefetch_multi = prefetch_multi
        wrapper.dirty_multi = dirty_multi

        return wrapper

    return decorator