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