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 hashlib
  12 import fcntl
  13
  14 class DiskCacheSlot(object):
  15     __slots__ = ("locator", "ready", "content", "cachedir")
  16
  17     def __init__(self, locator, cachedir):
  18         self.locator = locator
  19         self.ready = threading.Event()
  20         self.content = None
  21         self.cachedir = cachedir
  22
  23     def get(self):
  24         self.ready.wait()
  25         return self.content
  26
  27     def set(self, value):
  28         try:
  29             if value is None:
  30                 self.content = None
  31                 return
  32
  33             if len(value) == 0:
  34                 # Can't mmap a 0 length file
  35                 self.content = b''
  36                 return
  37
  38             if self.content is not None:
  39                 # Has been set already
  40                 return
  41
  42             blockdir = os.path.join(self.cachedir, self.locator[0:3])
  43             os.makedirs(blockdir, mode=0o700, exist_ok=True)
  44
  45             final = os.path.join(blockdir, self.locator)
  46
  47             f = tempfile.NamedTemporaryFile(dir=blockdir, delete=False)
  48             tmpfile = f.name
  49             os.chmod(tmpfile, stat.S_IRUSR | stat.S_IWUSR)
  50
  51             # aquire a shared lock, this tells other processes that
  52             # we're using this block and to please not delete it.
  53             fcntl.flock(f, fcntl.LOCK_SH)
  54
  55             f.write(value)
  56             f.flush()
  57             os.rename(tmpfile, final)
  58
  59             self.content = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
  60         except Exception as e:
  61             traceback.print_exc()
  62         finally:
  63             self.ready.set()
  64
  65     def size(self):
  66         if self.content is None:
  67             return 0
  68         else:
  69             return len(self.content)
  70
  71     def evict(self):
  72         if self.content is not None and len(self.content) > 0:
  73             # The mmap region might be in use when we decided to evict
  74             # it.  This can happen if the cache is too small.
  75             #
  76             # If we call close() now, it'll throw an error if
  77             # something tries to access it.
  78             #
  79             # However, we don't need to explicitly call mmap.close()
  80             #
  81             # I confirmed in mmapmodule.c that that both close
  82             # and deallocate do the same thing:
  83             #
  84             # a) close the file descriptor
  85             # b) unmap the memory range
  86             #
  87             # So we can forget it in the cache and delete the file on
  88             # disk, and it will tear it down after any other
  89             # lingering Python references to the mapped memory are
  90             # gone.
  91
  92             blockdir = os.path.join(self.cachedir, self.locator[0:3])
  93             final = os.path.join(blockdir, self.locator)
  94             try:
  95                 # If we can't upgrade our shared lock to an exclusive
  96                 # lock, it'll throw an error, that's fine and
  97                 # desirable, it means another process has a lock and
  98                 # we shouldn't delete the block.
  99                 fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
 100                 os.remove(final)
 101             except OSError:
 102                 pass
 103
 104     @staticmethod
 105     def get_from_disk(locator, cachedir):
 106         # Get it, check it, return it
 107         blockdir = os.path.join(cachedir, locator[0:3])
 108         final = os.path.join(blockdir, locator)
 109
 110         try:
 111             filehandle = open(final, "rb")
 112
 113             # aquire a shared lock, this tells other processes that
 114             # we're using this block and to please not delete it.
 115             fcntl.flock(f, fcntl.LOCK_SH)
 116
 117             content = mmap.mmap(filehandle.fileno(), 0, access=mmap.ACCESS_READ)
 118             disk_md5 = hashlib.md5(content).hexdigest()
 119             if disk_md5 == locator:
 120                 dc = DiskCacheSlot(locator, cachedir)
 121                 dc.content = content
 122                 dc.ready.set()
 123                 return dc
 124         except FileNotFoundError:
 125             pass
 126         except Exception as e:
 127             traceback.print_exc()
 128
 129         return None
 130
 131     @staticmethod
 132     def cleanup_cachedir(cachedir, maxsize):
 133         blocks = []
 134         totalsize = 0
 135         for root, dirs, files in os.walk(cachedir):
 136             for name in files:
 137                 blockpath = os.path.join(root, name)
 138                 res = os.stat(blockpath)
 139                 blocks.append((blockpath, res.st_size, res.st_atime))
 140                 totalsize += res.st_size
 141
 142         if totalsize <= maxsize:
 143             return
 144
 145         # sort by atime, so the blocks accessed the longest time in
 146         # the past get deleted first.
 147         blocks.sort(key=lambda x: x[2])
 148
 149         # go through the list and try deleting blocks until we're
 150         # below the target size and/or we run out of blocks
 151         i = 0
 152         while i < len(blocks) and totalsize > maxsize:
 153             try:
 154                 with open(blocks[i][0], "rb") as f:
 155                     # If we can't get an exclusive lock, it'll
 156                     # throw an error, that's fine and desirable,
 157                     # it means another process has a lock and we
 158                     # shouldn't delete the block.
 159                     fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
 160                     os.remove(block)
 161                     totalsize -= blocks[i][1]
 162             except OSError:
 163                 pass
 164             i += 1