sdk/python/arvados/diskcache.py

   1 # Copyright (C) The Arvados Authors. All rights reserved.
   2 #
   3 # SPDX-License-Identifier: Apache-2.0
   4
   5 import threading
   6 import mmap
   7 import os
   8 import traceback
   9 import stat
  10 import tempfile
  11 import fcntl
  12 import time
  13 import errno
  14 import logging
  15
  16 _logger = logging.getLogger('arvados.keep')
  17
  18 cacheblock_suffix = ".keepcacheblock"
  19
  20 class DiskCacheSlot(object):
  21     __slots__ = ("locator", "ready", "content", "cachedir", "filehandle")
  22
  23     def __init__(self, locator, cachedir):
  24         self.locator = locator
  25         self.ready = threading.Event()
  26         self.content = None
  27         self.cachedir = cachedir
  28         self.filehandle = None
  29
  30     def get(self):
  31         self.ready.wait()
  32         return self.content
  33
  34     def set(self, value):
  35         tmpfile = None
  36         try:
  37             if value is None:
  38                 self.content = None
  39                 self.ready.set()
  40                 return
  41
  42             if len(value) == 0:
  43                 # Can't mmap a 0 length file
  44                 self.content = b''
  45                 self.ready.set()
  46                 return
  47
  48             if self.content is not None:
  49                 # Has been set already
  50                 self.ready.set()
  51                 return
  52
  53             blockdir = os.path.join(self.cachedir, self.locator[0:3])
  54             os.makedirs(blockdir, mode=0o700, exist_ok=True)
  55
  56             final = os.path.join(blockdir, self.locator) + cacheblock_suffix
  57
  58             self.filehandle = tempfile.NamedTemporaryFile(dir=blockdir, delete=False, prefix="tmp", suffix=cacheblock_suffix)
  59             tmpfile = self.filehandle.name
  60             os.chmod(tmpfile, stat.S_IRUSR | stat.S_IWUSR)
  61
  62             # aquire a shared lock, this tells other processes that
  63             # we're using this block and to please not delete it.
  64             fcntl.flock(self.filehandle, fcntl.LOCK_SH)
  65
  66             self.filehandle.write(value)
  67             self.filehandle.flush()
  68             os.rename(tmpfile, final)
  69             tmpfile = None
  70
  71             self.content = mmap.mmap(self.filehandle.fileno(), 0, access=mmap.ACCESS_READ)
  72             # only set the event when mmap is successful
  73             self.ready.set()
  74         finally:
  75             if tmpfile is not None:
  76                 # If the tempfile hasn't been renamed on disk yet, try to delete it.
  77                 try:
  78                     os.remove(tmpfile)
  79                 except:
  80                     pass
  81
  82     def size(self):
  83         if self.content is None:
  84             return 0
  85         else:
  86             return len(self.content)
  87
  88     def evict(self):
  89         if self.content is not None and len(self.content) > 0:
  90             # The mmap region might be in use when we decided to evict
  91             # it.  This can happen if the cache is too small.
  92             #
  93             # If we call close() now, it'll throw an error if
  94             # something tries to access it.
  95             #
  96             # However, we don't need to explicitly call mmap.close()
  97             #
  98             # I confirmed in mmapmodule.c that that both close
  99             # and deallocate do the same thing:
 100             #
 101             # a) close the file descriptor
 102             # b) unmap the memory range
 103             #
 104             # So we can forget it in the cache and delete the file on
 105             # disk, and it will tear it down after any other
 106             # lingering Python references to the mapped memory are
 107             # gone.
 108
 109             blockdir = os.path.join(self.cachedir, self.locator[0:3])
 110             final = os.path.join(blockdir, self.locator) + cacheblock_suffix
 111             try:
 112                 fcntl.flock(self.filehandle, fcntl.LOCK_UN)
 113
 114                 # try to get an exclusive lock, this ensures other
 115                 # processes are not using the block.  It is
 116                 # nonblocking and will throw an exception if we
 117                 # can't get it, which is fine because that means
 118                 # we just won't try to delete it.
 119                 #
 120                 # I should note here, the file locking is not
 121                 # strictly necessary, we could just remove it and
 122                 # the kernel would ensure that the underlying
 123                 # inode remains available as long as other
 124                 # processes still have the file open.  However, if
 125                 # you have multiple processes sharing the cache
 126                 # and deleting each other's files, you'll end up
 127                 # with a bunch of ghost files that don't show up
 128                 # in the file system but are still taking up
 129                 # space, which isn't particularly user friendly.
 130                 # The locking strategy ensures that cache blocks
 131                 # in use remain visible.
 132                 #
 133                 fcntl.flock(self.filehandle, fcntl.LOCK_EX | fcntl.LOCK_NB)
 134
 135                 os.remove(final)
 136                 return True
 137             except OSError:
 138                 pass
 139             finally:
 140                 self.filehandle = None
 141                 self.content = None
 142             return False
 143
 144     @staticmethod
 145     def get_from_disk(locator, cachedir):
 146         blockdir = os.path.join(cachedir, locator[0:3])
 147         final = os.path.join(blockdir, locator) + cacheblock_suffix
 148
 149         try:
 150             filehandle = open(final, "rb")
 151
 152             # aquire a shared lock, this tells other processes that
 153             # we're using this block and to please not delete it.
 154             fcntl.flock(filehandle, fcntl.LOCK_SH)
 155
 156             content = mmap.mmap(filehandle.fileno(), 0, access=mmap.ACCESS_READ)
 157             dc = DiskCacheSlot(locator, cachedir)
 158             dc.filehandle = filehandle
 159             dc.content = content
 160             dc.ready.set()
 161             return dc
 162         except FileNotFoundError:
 163             pass
 164         except Exception as e:
 165             traceback.print_exc()
 166
 167         return None
 168
 169     @staticmethod
 170     def cache_usage(cachedir):
 171         usage = 0
 172         for root, dirs, files in os.walk(cachedir):
 173             for name in files:
 174                 if not name.endswith(cacheblock_suffix):
 175                     continue
 176
 177                 blockpath = os.path.join(root, name)
 178                 res = os.stat(blockpath)
 179                 usage += res.st_size
 180         return usage
 181
 182
 183     @staticmethod
 184     def init_cache(cachedir, maxslots):
 185         #
 186         # First check the disk cache works at all by creating a 1 byte cache entry
 187         #
 188         checkexists = DiskCacheSlot.get_from_disk('0cc175b9c0f1b6a831c399e269772661', cachedir)
 189         ds = DiskCacheSlot('0cc175b9c0f1b6a831c399e269772661', cachedir)
 190         ds.set(b'a')
 191         if checkexists is None:
 192             # Don't keep the test entry around unless it existed beforehand.
 193             ds.evict()
 194
 195         # map in all the files in the cache directory, up to max slots.
 196         # after max slots, try to delete the excess blocks.
 197         #
 198         # this gives the calling process ownership of all the blocks
 199
 200         blocks = []
 201         for root, dirs, files in os.walk(cachedir):
 202             for name in files:
 203                 if not name.endswith(cacheblock_suffix):
 204                     continue
 205
 206                 blockpath = os.path.join(root, name)
 207                 res = os.stat(blockpath)
 208
 209                 if len(name) == (32+len(cacheblock_suffix)) and not name.startswith("tmp"):
 210                     blocks.append((name[0:32], res.st_atime))
 211                 elif name.startswith("tmp") and ((time.time() - res.st_mtime) > 60):
 212                     # found a temporary file more than 1 minute old,
 213                     # try to delete it.
 214                     try:
 215                         os.remove(blockpath)
 216                     except:
 217                         pass
 218
 219         # sort by access time (atime), going from most recently
 220         # accessed (highest timestamp) to least recently accessed
 221         # (lowest timestamp).
 222         blocks.sort(key=lambda x: x[1], reverse=True)
 223
 224         # Map in all the files we found, up to maxslots, if we exceed
 225         # maxslots, start throwing things out.
 226         cachelist = []
 227         for b in blocks:
 228             got = DiskCacheSlot.get_from_disk(b[0], cachedir)
 229             if got is None:
 230                 continue
 231             if len(cachelist) < maxslots:
 232                 cachelist.append(got)
 233             else:
 234                 # we found more blocks than maxslots, try to
 235                 # throw it out of the cache.
 236                 got.evict()
 237
 238         return cachelist