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