1  import errno
   2  import hashlib
   3  import logging
   4  import os
   5  import platform
   6  import stat
   7  import struct
   8  import tempfile
   9  import typing as _t
  10  from contextlib import contextmanager
  11  from pathlib import Path
  12  from time import sleep
  13  from time import time
  14  
  15  from cachelib.base import BaseCache
  16  from cachelib.serializers import FileSystemSerializer
  17  
  18  
  19  def _lazy_md5(string: bytes = b"") -> _t.Any:
  20      """Don't access ``hashlib.md5`` until runtime. FIPS builds may not include
  21      md5, in which case the import and use as a default would fail before the
  22      developer can configure something else.
  23      """
  24      return hashlib.md5(string)
  25  
  26  
  27  class FileSystemCache(BaseCache):
  28      """A cache that stores the items on the file system.  This cache depends
  29      on being the only user of the `cache_dir`.  Make absolutely sure that
  30      nobody but this cache stores files there or otherwise the cache will
  31      randomly delete files therein.
  32  
  33      :param cache_dir: the directory where cache files are stored.
  34      :param threshold: the maximum number of items the cache stores before
  35                        it starts deleting some. A threshold value of 0
  36                        indicates no threshold.
  37      :param default_timeout: the default timeout that is used if no timeout is
  38                              specified on :meth:`~BaseCache.set`. A timeout of
  39                              0 indicates that the cache never expires.
  40      :param mode: the file mode wanted for the cache files, default 0600
  41      :param hash_method: Default hashlib.md5. The hash method used to
  42                          generate the filename for cached results.
  43                          Default is lazy loaded and can be overriden by
  44                          seeting  `_default_hash_method`
  45      """
  46  
  47      #: used for temporary files by the FileSystemCache
  48      _fs_transaction_suffix = ".__wz_cache"
  49      #: keep amount of files in a cache element
  50      _fs_count_file = "__wz_cache_count"
  51      #: default file name hashing method
  52      _default_hash_method = staticmethod(_lazy_md5)
  53  
  54      serializer = FileSystemSerializer()
  55  
  56      def __init__(
  57          self,
  58          cache_dir: str,
  59          threshold: int = 500,
  60          default_timeout: int = 300,
  61          mode: _t.Optional[int] = None,
  62          hash_method: _t.Any = None,
  63      ):
  64          BaseCache.__init__(self, default_timeout)
  65          self._path = cache_dir
  66          self._threshold = threshold
  67  
  68          self._hash_method = self._default_hash_method
  69          if hash_method is not None:
  70              self._hash_method = hash_method
  71  
  72          # Mode set by user takes precedence. If no mode has
  73          # been given, we need to set the correct default based
  74          # on user platform.
  75          self._mode = mode
  76          if self._mode is None:
  77              self._mode = self._get_compatible_platform_mode()
  78  
  79          try:
  80              os.makedirs(self._path)
  81          except OSError as ex:
  82              if ex.errno != errno.EEXIST:
  83                  raise
  84  
  85          # If there are many files and a zero threshold,
  86          # the list_dir can slow initialisation massively
  87          if self._threshold != 0:
  88              self._update_count(value=len(list(self._list_dir())))
  89  
  90      def _get_compatible_platform_mode(self) -> int:
  91          mode = 0o600  # nix systems
  92          if platform.system() == "Windows":
  93              mode = stat.S_IWRITE
  94          return mode
  95  
  96      @property
  97      def _file_count(self) -> int:
  98          return self.get(self._fs_count_file) or 0
  99  
 100      def _update_count(
 101          self, delta: _t.Optional[int] = None, value: _t.Optional[int] = None
 102      ) -> None:
 103          # If we have no threshold, don't count files
 104          if self._threshold == 0:
 105              return
 106          if delta:
 107              new_count = self._file_count + delta
 108          else:
 109              new_count = value or 0
 110          self.set(self._fs_count_file, new_count, mgmt_element=True)
 111  
 112      def _normalize_timeout(self, timeout: _t.Optional[int]) -> int:
 113          timeout = BaseCache._normalize_timeout(self, timeout)
 114          if timeout != 0:
 115              timeout = int(time()) + timeout
 116          return int(timeout)
 117  
 118      def _is_mgmt(self, name: str) -> bool:
 119          fshash = self._get_filename(self._fs_count_file).split(os.sep)[-1]
 120          return name == fshash or name.endswith(self._fs_transaction_suffix)
 121  
 122      def _list_dir(self) -> _t.Generator[str, None, None]:
 123          """return a list of (fully qualified) cache filenames"""
 124          return (
 125              os.path.join(self._path, fn)
 126              for fn in os.listdir(self._path)
 127              if not self._is_mgmt(fn)
 128          )
 129  
 130      def _over_threshold(self) -> bool:
 131          return self._threshold != 0 and self._file_count > self._threshold
 132  
 133      def _remove_expired(self, now: float) -> None:
 134          for fname in self._list_dir():
 135              try:
 136                  with self._safe_stream_open(fname, "rb") as f:
 137                      expires = struct.unpack("I", f.read(4))[0]
 138                  if expires != 0 and expires < now:
 139                      os.remove(fname)
 140                      self._update_count(delta=-1)
 141              except FileNotFoundError:
 142                  pass
 143              except (OSError, EOFError, struct.error):
 144                  logging.warning(
 145                      "Exception raised while handling cache file '%s'",
 146                      fname,
 147                      exc_info=True,
 148                  )
 149  
 150      def _remove_older(self) -> bool:
 151          exp_fname_tuples = []
 152          for fname in self._list_dir():
 153              try:
 154                  with self._safe_stream_open(fname, "rb") as f:
 155                      timestamp = struct.unpack("I", f.read(4))[0]
 156                      exp_fname_tuples.append((timestamp, fname))
 157              except FileNotFoundError:
 158                  pass
 159              except (OSError, EOFError, struct.error):
 160                  logging.warning(
 161                      "Exception raised while handling cache file '%s'",
 162                      fname,
 163                      exc_info=True,
 164                  )
 165          fname_sorted = (
 166              fname for _, fname in sorted(exp_fname_tuples, key=lambda item: item[0])
 167          )
 168          for fname in fname_sorted:
 169              try:
 170                  os.remove(fname)
 171                  self._update_count(delta=-1)
 172              except FileNotFoundError:
 173                  pass
 174              except OSError:
 175                  logging.warning(
 176                      "Exception raised while handling cache file '%s'",
 177                      fname,
 178                      exc_info=True,
 179                  )
 180                  return False
 181              if not self._over_threshold():
 182                  break
 183          return True
 184  
 185      def _prune(self) -> None:
 186          if self._over_threshold():
 187              now = time()
 188              self._remove_expired(now)
 189          # if still over threshold
 190          if self._over_threshold():
 191              self._remove_older()
 192  
 193      def clear(self) -> bool:
 194          for i, fname in enumerate(self._list_dir()):
 195              try:
 196                  os.remove(fname)
 197              except FileNotFoundError:
 198                  pass
 199              except OSError:
 200                  logging.warning(
 201                      "Exception raised while handling cache file '%s'",
 202                      fname,
 203                      exc_info=True,
 204                  )
 205                  self._update_count(delta=-i)
 206                  return False
 207          self._update_count(value=0)
 208          return True
 209  
 210      def _get_filename(self, key: str) -> str:
 211          if isinstance(key, str):
 212              bkey = key.encode("utf-8")  # XXX unicode review
 213              bkey_hash = self._hash_method(bkey).hexdigest()
 214          else:
 215              raise TypeError(f"Key must be a string, received type {type(key)}")
 216          return os.path.join(self._path, bkey_hash)
 217  
 218      def get(self, key: str) -> _t.Any:
 219          filename = self._get_filename(key)
 220          try:
 221              with self._safe_stream_open(filename, "rb") as f:
 222                  pickle_time = struct.unpack("I", f.read(4))[0]
 223                  if pickle_time == 0 or pickle_time >= time():
 224                      return self.serializer.load(f)
 225          except FileNotFoundError:
 226              pass
 227          except (OSError, EOFError, struct.error):
 228              logging.warning(
 229                  "Exception raised while handling cache file '%s'",
 230                  filename,
 231                  exc_info=True,
 232              )
 233          return None
 234  
 235      def add(self, key: str, value: _t.Any, timeout: _t.Optional[int] = None) -> bool:
 236          filename = self._get_filename(key)
 237          if not os.path.exists(filename):
 238              return self.set(key, value, timeout)
 239          return False
 240  
 241      def set(
 242          self,
 243          key: str,
 244          value: _t.Any,
 245          timeout: _t.Optional[int] = None,
 246          mgmt_element: bool = False,
 247      ) -> bool:
 248          # Management elements have no timeout
 249          if mgmt_element:
 250              timeout = 0
 251          # Don't prune on management element update, to avoid loop
 252          else:
 253              self._prune()
 254  
 255          timeout = self._normalize_timeout(timeout)
 256          filename = self._get_filename(key)
 257          overwrite = os.path.isfile(filename)
 258  
 259          try:
 260              fd, tmp = tempfile.mkstemp(
 261                  suffix=self._fs_transaction_suffix, dir=self._path
 262              )
 263              with os.fdopen(fd, "wb") as f:
 264                  f.write(struct.pack("I", timeout))
 265                  self.serializer.dump(value, f)
 266  
 267              self._run_safely(os.replace, tmp, filename)
 268              self._run_safely(os.chmod, filename, self._mode)
 269  
 270              fsize = Path(filename).stat().st_size
 271          except OSError:
 272              logging.warning(
 273                  "Exception raised while handling cache file '%s'",
 274                  filename,
 275                  exc_info=True,
 276              )
 277              return False
 278          else:
 279              # Management elements should not count towards threshold
 280              if not overwrite and not mgmt_element:
 281                  self._update_count(delta=1)
 282              return fsize > 0  # function should fail if file is empty
 283  
 284      def delete(self, key: str, mgmt_element: bool = False) -> bool:
 285          try:
 286              os.remove(self._get_filename(key))
 287          except FileNotFoundError:  # if file doesn't exist we consider it deleted
 288              return True
 289          except OSError:
 290              logging.warning("Exception raised while handling cache file", exc_info=True)
 291              return False
 292          else:
 293              # Management elements should not count towards threshold
 294              if not mgmt_element:
 295                  self._update_count(delta=-1)
 296              return True
 297  
 298      def has(self, key: str) -> bool:
 299          filename = self._get_filename(key)
 300          try:
 301              with self._safe_stream_open(filename, "rb") as f:
 302                  pickle_time = struct.unpack("I", f.read(4))[0]
 303                  if pickle_time == 0 or pickle_time >= time():
 304                      return True
 305                  else:
 306                      return False
 307          except FileNotFoundError:  # if there is no file there is no key
 308              return False
 309          except (OSError, EOFError, struct.error):
 310              logging.warning(
 311                  "Exception raised while handling cache file '%s'",
 312                  filename,
 313                  exc_info=True,
 314              )
 315              return False
 316  
 317      def _run_safely(self, fn: _t.Callable, *args: _t.Any, **kwargs: _t.Any) -> _t.Any:
 318          """On Windows os.replace, os.chmod and open can yield
 319          permission errors if executed by two different processes."""
 320          if platform.system() == "Windows":
 321              output = None
 322              wait_step = 0.001
 323              max_sleep_time = 10.0
 324              total_sleep_time = 0.0
 325  
 326              while total_sleep_time < max_sleep_time:
 327                  try:
 328                      output = fn(*args, **kwargs)
 329                  except PermissionError:
 330                      sleep(wait_step)
 331                      total_sleep_time += wait_step
 332                      wait_step *= 2
 333                  else:
 334                      break
 335          else:
 336              output = fn(*args, **kwargs)
 337  
 338          return output
 339  
 340      @contextmanager
 341      def _safe_stream_open(self, path: str, mode: str) -> _t.Generator:
 342          fs = self._run_safely(open, path, mode)
 343          if fs is None:
 344              raise OSError
 345          try:
 346              yield fs
 347          finally:
 348              fs.close()
 349