bt-source/panel/class/cachelib/file.py

import errno
import logging
import os
import platform
import stat
import struct
import tempfile
import typing as _t
from contextlib import contextmanager
from hashlib import md5
from pathlib import Path
from time import sleep
from time import time

from cachelib.base import BaseCache
from cachelib.serializers import FileSystemSerializer
from cachelib.simple import SimpleCache


class FileSystemCache(BaseCache):
    """A cache that stores the items on the file system.  This cache depends
    on being the only user of the `cache_dir`.  Make absolutely sure that
    nobody but this cache stores files there or otherwise the cache will
    randomly delete files therein.

    :param cache_dir: the directory where cache files are stored.
    :param threshold: the maximum number of items the cache stores before
                      it starts deleting some. A threshold value of 0
                      indicates no threshold.
    :param default_timeout: the default timeout that is used if no timeout is
                            specified on :meth:`~BaseCache.set`. A timeout of
                            0 indicates that the cache never expires.
    :param mode: the file mode wanted for the cache files, default 0600
    :param hash_method: Default hashlib.md5. The hash method used to
                        generate the filename for cached results.
    """

    #: used for temporary files by the FileSystemCache
    _fs_transaction_suffix = ".__wz_cache"
    #: keep amount of files in a cache element
    _fs_count_file = "__wz_cache_count"

    serializer = FileSystemSerializer()

    def __init__(
        self,
        cache_dir: str,
        threshold: int = 500,
        default_timeout: int = 300,
        mode: _t.Optional[int] = None,
        hash_method: _t.Any = md5,
    ):
        BaseCache.__init__(self, default_timeout)
        self._path = cache_dir
        self._threshold = threshold
        self._hash_method = hash_method
        # 当出现磁盘空间不足时，触发使用内存进行会话存储的能力，保证用户可登录，操作删除垃圾文件
        self._mem_cache = SimpleCache(threshold=threshold, default_timeout=default_timeout)

        # Mode set by user takes precedence. If no mode has
        # been given, we need to set the correct default based
        # on user platform.
        self._mode = mode
        if self._mode is None:
            self._mode = self._get_compatible_platform_mode()

        try:
            os.makedirs(self._path)
        except OSError as ex:
            if ex.errno != errno.EEXIST:
                raise

        # If there are many files and a zero threshold,
        # the list_dir can slow initialisation massively
        if self._threshold != 0:
            self._update_count(value=len(list(self._list_dir())))

    def _get_compatible_platform_mode(self) -> int:
        mode = 0o600  # nix systems
        if platform.system() == "Windows":
            mode = stat.S_IWRITE
        return mode

    @property
    def _file_count(self) -> int:
        return self.get(self._fs_count_file) or 0

    def _update_count(
        self, delta: _t.Optional[int] = None, value: _t.Optional[int] = None
    ) -> None:
        # If we have no threshold, don't count files
        if self._threshold == 0:
            return
        if delta:
            new_count = self._file_count + delta
        else:
            new_count = value or 0
        self.set(self._fs_count_file, new_count, mgmt_element=True)

    def _normalize_timeout(self, timeout: _t.Optional[int]) -> int:
        timeout = BaseCache._normalize_timeout(self, timeout)
        if timeout != 0:
            timeout = int(time()) + timeout
        return int(timeout)

    def _is_mgmt(self, name: str) -> bool:
        fshash = self._get_filename(self._fs_count_file).split(os.sep)[-1]
        return name == fshash or name.endswith(self._fs_transaction_suffix)

    def _list_dir(self) -> _t.Generator[str, None, None]:
        """return a list of (fully qualified) cache filenames"""
        return (
            os.path.join(self._path, fn)
            for fn in os.listdir(self._path)
            if not self._is_mgmt(fn)
        )

    def _over_threshold(self) -> bool:
        return self._threshold != 0 and self._file_count > self._threshold

    def _remove_expired(self, now: float) -> None:
        for fname in self._list_dir():
            try:
                with self._safe_stream_open(fname, "rb") as f:
                    expires = struct.unpack("I", f.read(4))[0]
                if expires != 0 and expires < now:
                    os.remove(fname)
                    self._update_count(delta=-1)
            except FileNotFoundError:
                pass
            except (OSError, EOFError, struct.error):
                logging.warning(
                    "Exception raised while handling cache file '%s'",
                    fname,
                    exc_info=True,
                )

    def _remove_older(self) -> bool:
        exp_fname_tuples = []
        for fname in self._list_dir():
            try:
                with self._safe_stream_open(fname, "rb") as f:
                    timestamp = struct.unpack("I", f.read(4))[0]
                    exp_fname_tuples.append((timestamp, fname))
            except FileNotFoundError:
                pass
            except (OSError, EOFError, struct.error):
                logging.warning(
                    "Exception raised while handling cache file '%s'",
                    fname,
                    exc_info=True,
                )
        fname_sorted = (
            fname
            for _, fname in sorted(
                exp_fname_tuples, key=lambda item: item[0]  # type: ignore
            )
        )
        for fname in fname_sorted:
            try:
                os.remove(fname)
                self._update_count(delta=-1)
            except FileNotFoundError:
                pass
            except OSError:
                logging.warning(
                    "Exception raised while handling cache file '%s'",
                    fname,
                    exc_info=True,
                )
                return False
            if not self._over_threshold():
                break
        return True

    def _prune(self) -> None:
        if self._over_threshold():
            now = time()
            self._remove_expired(now)
        # if still over threshold
        if self._over_threshold():
            self._remove_older()

    def clear(self) -> bool:
        for i, fname in enumerate(self._list_dir()):
            try:
                os.remove(fname)
            except FileNotFoundError:
                pass
            except OSError:
                logging.warning(
                    "Exception raised while handling cache file '%s'",
                    fname,
                    exc_info=True,
                )
                self._update_count(delta=-i)
                return False
        self._update_count(value=0)
        self._mem_cache.clear()
        return True

    def _get_filename(self, key: str) -> str:
        if isinstance(key, str):
            bkey = key.encode("utf-8")  # XXX unicode review
            bkey_hash = self._hash_method(bkey).hexdigest()
        else:
            raise TypeError(f"Key must be a string, received type {type(key)}")
        return os.path.join(self._path, bkey_hash)

    def get(self, key: str) -> _t.Any:
        if not isinstance(key, str): return False
        filename = self._get_filename(key)
        try:
            with self._safe_stream_open(filename, "rb") as f:
                pickle_time = struct.unpack("I", f.read(4))[0]
                if pickle_time == 0 or pickle_time >= time():
                    return self.serializer.load(f)
        except FileNotFoundError:
            pass
        except (OSError, EOFError, struct.error):
            logging.warning(
                "Exception raised while handling cache file '%s'",
                filename,
                exc_info=True,
            )
        if self._mem_cache.has(key):
            return self._mem_cache.get(key)
        return None

    def add(self, key: str, value: _t.Any, timeout: _t.Optional[int] = None) -> bool:
        if not isinstance(key, str): return False
        type_list = (int, float, bool, str, list, dict, tuple, set, bytes)
        value_type = type(value)
        if value_type not in type_list:
            return False
        filename = self._get_filename(key)
        if not os.path.exists(filename):
            return self.set(key, value, timeout)
        return False

    def set(
        self,
        key: str,
        value: _t.Any,
        timeout: _t.Optional[int] = None,
        mgmt_element: bool = False,
    ) -> bool:
        if not isinstance(key, str): return False
        type_list = (int, float, bool, str, list, dict, tuple, set, bytes)
        value_type = type(value)
        if value_type not in type_list:
            return False
        # Management elements have no timeout
        if mgmt_element:
            timeout = 0
        # Don't prune on management element update, to avoid loop
        else:
            self._prune()

        timeout = self._normalize_timeout(timeout)
        filename = self._get_filename(key)
        overwrite = os.path.isfile(filename)

        try:
            fd, tmp = tempfile.mkstemp(
                suffix=self._fs_transaction_suffix, dir=self._path
            )
            with os.fdopen(fd, "wb") as f:
                f.write(struct.pack("I", timeout))
                self.serializer.dump(value, f)

            self._run_safely(os.replace, tmp, filename)
            self._run_safely(os.chmod, filename, self._mode)

            fsize = Path(filename).stat().st_size
        except OSError as e:
            if "No space left on device" in str(e):
                return self._mem_cache.set(key, value, timeout)
            logging.warning(
                "Exception raised while handling cache file '%s'",
                filename,
                exc_info=True,
            )
            return False
        else:
            # Management elements should not count towards threshold
            if not overwrite and not mgmt_element:
                self._update_count(delta=1)
            return fsize > 0  # function should fail if file is empty

    def delete(self, key: str, mgmt_element: bool = False) -> bool:
        if self._mem_cache.has(key):
            self._mem_cache.delete(key)
        try:
            os.remove(self._get_filename(key))
        except FileNotFoundError:  # if file doesn't exist we consider it deleted
            return True
        except OSError:
            logging.warning("Exception raised while handling cache file", exc_info=True)
            return False
        else:
            # Management elements should not count towards threshold
            if not mgmt_element:
                self._update_count(delta=-1)
            return True

    def has(self, key: str) -> bool:
        if self._mem_cache.has(key):
            return True
        filename = self._get_filename(key)
        try:
            with self._safe_stream_open(filename, "rb") as f:
                pickle_time = struct.unpack("I", f.read(4))[0]
                if pickle_time == 0 or pickle_time >= time():
                    return True
                else:
                    return False
        except FileNotFoundError:  # if there is no file there is no key
            return False
        except (OSError, EOFError, struct.error):
            logging.warning(
                "Exception raised while handling cache file '%s'",
                filename,
                exc_info=True,
            )
            return False

    def _run_safely(self, fn: _t.Callable, *args: _t.Any, **kwargs: _t.Any) -> _t.Any:
        """On Windows os.replace, os.chmod and open can yield
        permission errors if executed by two different processes."""
        if platform.system() == "Windows":
            output = None
            wait_step = 0.001
            max_sleep_time = 10.0
            total_sleep_time = 0.0

            while total_sleep_time < max_sleep_time:
                try:
                    output = fn(*args, **kwargs)
                except PermissionError:
                    sleep(wait_step)
                    total_sleep_time += wait_step
                    wait_step *= 2
                else:
                    break
        else:
            output = fn(*args, **kwargs)

        return output

    @contextmanager
    def _safe_stream_open(self, path: str, mode: str) -> _t.Generator:
        fs = self._run_safely(open, path, mode)
        if fs is None:
            raise OSError
        try:
            yield fs
        finally:
            fs.close()