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
  12 class DiskCacheSlot(object):
  13     __slots__ = ("locator", "ready", "content", "cachedir")
  14
  15     def __init__(self, locator, cachedir):
  16         self.locator = locator
  17         self.ready = threading.Event()
  18         self.content = None
  19         self.cachedir = cachedir
  20
  21     def get(self):
  22         self.ready.wait()
  23         return self.content
  24
  25     def set(self, value):
  26         try:
  27             if value is None:
  28                 self.content = None
  29                 return
  30
  31             if len(value) == 0:
  32                 # Can't mmap a 0 length file
  33                 self.content = b''
  34                 return
  35
  36             if self.content is not None:
  37                 # Has been set already
  38                 return
  39
  40             blockdir = os.path.join(self.cachedir, self.locator[0:3])
  41             os.makedirs(blockdir, mode=0o700, exist_ok=True)
  42
  43             final = os.path.join(blockdir, self.locator)
  44
  45             f = tempfile.NamedTemporaryFile(dir=blockdir, delete=False)
  46             tmpfile = f.name
  47             os.chmod(tmpfile, stat.S_IRUSR | stat.S_IWUSR)
  48
  49             f.write(value)
  50             f.flush()
  51             os.rename(tmpfile, final)
  52
  53             self.content = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
  54         except Exception as e:
  55             traceback.print_exc()
  56         finally:
  57             self.ready.set()
  58
  59     def size(self):
  60         if self.content is None:
  61             return 0
  62         else:
  63             return len(self.content)
  64
  65     def evict(self):
  66         if self.content is not None and len(self.content) > 0:
  67             # The mmap region might be in use when we decided to evict
  68             # it.  This can happen if the cache is too small.
  69             #
  70             # If we call close() now, it'll throw an error if
  71             # something tries to access it.
  72             #
  73             # However, we don't need to explicitly call mmap.close()
  74             #
  75             # I confirmed in mmapmodule.c that that both close
  76             # and deallocate do the same thing:
  77             #
  78             # a) close the file descriptor
  79             # b) unmap the memory range
  80             #
  81             # So we can forget it in the cache and delete the file on
  82             # disk, and it will tear it down after any other
  83             # lingering Python references to the mapped memory are
  84             # gone.
  85
  86             blockdir = os.path.join(self.cachedir, self.locator[0:3])
  87             final = os.path.join(blockdir, self.locator)
  88             try:
  89                 os.remove(final)
  90             except OSError:
  91                 pass
  92
  93     @staticmethod
  94     def get_from_disk(locator, cachedir):
  95         blockdir = os.path.join(cachedir, locator[0:3])
  96         final = os.path.join(blockdir, locator)
  97
  98         try:
  99             filehandle = open(final, "rb")
 100
 101             content = mmap.mmap(filehandle.fileno(), 0, access=mmap.ACCESS_READ)
 102             dc = DiskCacheSlot(locator, cachedir)
 103             dc.content = content
 104             dc.ready.set()
 105             return dc
 106         except FileNotFoundError:
 107             pass
 108         except Exception as e:
 109             traceback.print_exc()
 110
 111         return None
 112
 113     @staticmethod
 114     def init_cache(cachedir, maxslots):
 115         # map in all the files in the cache directory, up to max slots.
 116         # after max slots, try to delete the excess blocks.
 117         #
 118         # this gives the calling process ownership of all the blocks
 119
 120         blocks = []
 121         for root, dirs, files in os.walk(cachedir):
 122             for name in files:
 123                 blockpath = os.path.join(root, name)
 124                 res = os.stat(blockpath)
 125                 blocks.append((name, res.st_atime))
 126
 127         # sort by access time (atime), going from most recently
 128         # accessed (highest timestamp) to least recently accessed
 129         # (lowest timestamp).
 130         blocks.sort(key=lambda x: x[1], reverse=True)
 131
 132         # Map in all the files we found, up to maxslots, if we exceed
 133         # maxslots, start throwing things out.
 134         cachelist = []
 135         for b in blocks:
 136             got = DiskCacheSlot.get_from_disk(b[0], cachedir)
 137             if got is None:
 138                 continue
 139             if len(cachelist) < maxslots:
 140                 cachelist.append(got)
 141             else:
 142                 # we found more blocks than maxslots, try to
 143                 # throw it out of the cache.
 144                 got.evict()
 145
 146         return cachelist