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
  13 class DiskCacheSlot(object):
  14     __slots__ = ("locator", "ready", "content", "cachedir")
  15
  16     def __init__(self, locator, cachedir):
  17         self.locator = locator
  18         self.ready = threading.Event()
  19         self.content = None
  20         self.cachedir = cachedir
  21
  22     def get(self):
  23         self.ready.wait()
  24         return self.content
  25
  26     def set(self, value):
  27         try:
  28             if value is None:
  29                 self.content = None
  30                 return
  31
  32             if len(value) == 0:
  33                 # Can't mmap a 0 length file
  34                 self.content = b''
  35                 return
  36
  37             if self.content is not None:
  38                 # Has been set already
  39                 return
  40
  41             blockdir = os.path.join(self.cachedir, self.locator[0:3])
  42             os.makedirs(blockdir, mode=0o700, exist_ok=True)
  43
  44             final = os.path.join(blockdir, self.locator)
  45
  46             f = tempfile.NamedTemporaryFile(dir=blockdir, delete=False)
  47             tmpfile = f.name
  48             os.chmod(tmpfile, stat.S_IRUSR | stat.S_IWUSR)
  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         # Get it, check it, return it
  96         blockdir = os.path.join(cachedir, locator[0:3])
  97         final = os.path.join(blockdir, locator)
  98
  99         try:
 100             filehandle = open(final, "rb")
 101             content = mmap.mmap(filehandle.fileno(), 0, access=mmap.ACCESS_READ)
 102             disk_md5 = hashlib.md5(content).hexdigest()
 103             if disk_md5 == locator:
 104                 dc = DiskCacheSlot(locator, cachedir)
 105                 dc.content = content
 106                 dc.ready.set()
 107                 return dc
 108         except FileNotFoundError:
 109             pass
 110         except Exception as e:
 111             traceback.print_exc()
 112
 113         return None