Merge branch '21639-keep-cache-dict' refs #21639
[arvados.git] / sdk / python / arvados / diskcache.py
index 9734d93a7742d6bf395b4fe2dbed2e907ab50a4d..528a7d28b58146af1a33eac0ada4b746a9eaa12d 100644 (file)
@@ -9,24 +9,38 @@ import traceback
 import stat
 import tempfile
 import fcntl
+import time
 import errno
 import logging
+import weakref
+import collections
 
 _logger = logging.getLogger('arvados.keep')
 
 cacheblock_suffix = ".keepcacheblock"
 
 class DiskCacheSlot(object):
-    __slots__ = ("locator", "ready", "content", "cachedir")
+    __slots__ = ("locator", "ready", "content", "cachedir", "filehandle", "linger")
 
     def __init__(self, locator, cachedir):
         self.locator = locator
         self.ready = threading.Event()
         self.content = None
         self.cachedir = cachedir
+        self.filehandle = None
+        self.linger = None
 
     def get(self):
         self.ready.wait()
+        # 'content' can None, an empty byte string, or a nonempty mmap
+        # region.  If it is an mmap region, we want to advise the
+        # kernel we're going to use it.  This nudges the kernel to
+        # re-read most or all of the block if necessary (instead of
+        # just a few pages at a time), reducing the number of page
+        # faults and improving performance by 4x compared to not
+        # calling madvise.
+        if self.content:
+            self.content.madvise(mmap.MADV_WILLNEED)
         return self.content
 
     def set(self, value):
@@ -34,47 +48,42 @@ class DiskCacheSlot(object):
         try:
             if value is None:
                 self.content = None
-                return
+                self.ready.set()
+                return False
 
             if len(value) == 0:
                 # Can't mmap a 0 length file
                 self.content = b''
-                return
+                self.ready.set()
+                return True
 
             if self.content is not None:
                 # Has been set already
-                return
+                self.ready.set()
+                return False
 
             blockdir = os.path.join(self.cachedir, self.locator[0:3])
             os.makedirs(blockdir, mode=0o700, exist_ok=True)
 
             final = os.path.join(blockdir, self.locator) + cacheblock_suffix
 
-            f = tempfile.NamedTemporaryFile(dir=blockdir, delete=False, prefix="tmp", suffix=cacheblock_suffix)
-            tmpfile = f.name
+            self.filehandle = tempfile.NamedTemporaryFile(dir=blockdir, delete=False, prefix="tmp", suffix=cacheblock_suffix)
+            tmpfile = self.filehandle.name
             os.chmod(tmpfile, stat.S_IRUSR | stat.S_IWUSR)
 
             # aquire a shared lock, this tells other processes that
             # we're using this block and to please not delete it.
-            fcntl.flock(f, fcntl.LOCK_SH)
+            fcntl.flock(self.filehandle, fcntl.LOCK_SH)
 
-            f.write(value)
-            f.flush()
+            self.filehandle.write(value)
+            self.filehandle.flush()
             os.rename(tmpfile, final)
             tmpfile = None
 
-            self.content = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
-        except OSError as e:
-            if e.errno == errno.ENODEV:
-                _logger.error("Unable to use disk cache: The underlying filesystem does not support memory mapping.")
-            elif e.errno == errno.ENOMEM:
-                _logger.error("Unable to use disk cache: The process's maximum number of mappings would have been exceeded.")
-            elif e.errno == errno.ENOSPC:
-                _logger.error("Unable to use disk cache: Out of disk space.")
-            else:
-                traceback.print_exc()
-        except Exception as e:
-            traceback.print_exc()
+            self.content = mmap.mmap(self.filehandle.fileno(), 0, access=mmap.ACCESS_READ)
+            # only set the event when mmap is successful
+            self.ready.set()
+            return True
         finally:
             if tmpfile is not None:
                 # If the tempfile hasn't been renamed on disk yet, try to delete it.
@@ -82,74 +91,76 @@ class DiskCacheSlot(object):
                     os.remove(tmpfile)
                 except:
                     pass
-            if self.content is None:
-                # Something went wrong with the disk cache, fall back
-                # to RAM cache behavior (the alternative is to cache
-                # nothing and return a read error).
-                self.content = value
-            self.ready.set()
 
     def size(self):
         if self.content is None:
+            if self.linger is not None:
+                # If it is still lingering (object is still accessible
+                # through the weak reference) it is still taking up
+                # space.
+                content = self.linger()
+                if content is not None:
+                    return len(content)
             return 0
         else:
             return len(self.content)
 
     def evict(self):
-        if self.content is not None and len(self.content) > 0:
-            # The mmap region might be in use when we decided to evict
-            # it.  This can happen if the cache is too small.
-            #
-            # If we call close() now, it'll throw an error if
-            # something tries to access it.
-            #
-            # However, we don't need to explicitly call mmap.close()
-            #
-            # I confirmed in mmapmodule.c that that both close
-            # and deallocate do the same thing:
+        if not self.content:
+            return
+
+        # The mmap region might be in use when we decided to evict
+        # it.  This can happen if the cache is too small.
+        #
+        # If we call close() now, it'll throw an error if
+        # something tries to access it.
+        #
+        # However, we don't need to explicitly call mmap.close()
+        #
+        # I confirmed in mmapmodule.c that that both close
+        # and deallocate do the same thing:
+        #
+        # a) close the file descriptor
+        # b) unmap the memory range
+        #
+        # So we can forget it in the cache and delete the file on
+        # disk, and it will tear it down after any other
+        # lingering Python references to the mapped memory are
+        # gone.
+
+        blockdir = os.path.join(self.cachedir, self.locator[0:3])
+        final = os.path.join(blockdir, self.locator) + cacheblock_suffix
+        try:
+            fcntl.flock(self.filehandle, fcntl.LOCK_UN)
+
+            # try to get an exclusive lock, this ensures other
+            # processes are not using the block.  It is
+            # nonblocking and will throw an exception if we
+            # can't get it, which is fine because that means
+            # we just won't try to delete it.
             #
-            # a) close the file descriptor
-            # b) unmap the memory range
+            # I should note here, the file locking is not
+            # strictly necessary, we could just remove it and
+            # the kernel would ensure that the underlying
+            # inode remains available as long as other
+            # processes still have the file open.  However, if
+            # you have multiple processes sharing the cache
+            # and deleting each other's files, you'll end up
+            # with a bunch of ghost files that don't show up
+            # in the file system but are still taking up
+            # space, which isn't particularly user friendly.
+            # The locking strategy ensures that cache blocks
+            # in use remain visible.
             #
-            # So we can forget it in the cache and delete the file on
-            # disk, and it will tear it down after any other
-            # lingering Python references to the mapped memory are
-            # gone.
+            fcntl.flock(self.filehandle, fcntl.LOCK_EX | fcntl.LOCK_NB)
 
-            blockdir = os.path.join(self.cachedir, self.locator[0:3])
-            final = os.path.join(blockdir, self.locator) + cacheblock_suffix
-            try:
-                with open(final, "rb") as f:
-                    # unlock
-                    fcntl.flock(f, fcntl.LOCK_UN)
-                    self.content = None
-
-                    # try to get an exclusive lock, this ensures other
-                    # processes are not using the block.  It is
-                    # nonblocking and will throw an exception if we
-                    # can't get it, which is fine because that means
-                    # we just won't try to delete it.
-                    #
-                    # I should note here, the file locking is not
-                    # strictly necessary, we could just remove it and
-                    # the kernel would ensure that the underlying
-                    # inode remains available as long as other
-                    # processes still have the file open.  However, if
-                    # you have multiple processes sharing the cache
-                    # and deleting each other's files, you'll end up
-                    # with a bunch of ghost files that don't show up
-                    # in the file system but are still taking up
-                    # space, which isn't particularly user friendly.
-                    # The locking strategy ensures that cache blocks
-                    # in use remain visible.
-                    #
-                    fcntl.flock(filehandle, fcntl.LOCK_EX | fcntl.LOCK_NB)
-
-                    os.remove(final)
-                    return True
-            except OSError:
-                pass
-            return False
+            os.remove(final)
+            return True
+        except OSError:
+            pass
+        finally:
+            self.filehandle = None
+            self.content = None
 
     @staticmethod
     def get_from_disk(locator, cachedir):
@@ -165,6 +176,7 @@ class DiskCacheSlot(object):
 
             content = mmap.mmap(filehandle.fileno(), 0, access=mmap.ACCESS_READ)
             dc = DiskCacheSlot(locator, cachedir)
+            dc.filehandle = filehandle
             dc.content = content
             dc.ready.set()
             return dc
@@ -175,8 +187,32 @@ class DiskCacheSlot(object):
 
         return None
 
+    @staticmethod
+    def cache_usage(cachedir):
+        usage = 0
+        for root, dirs, files in os.walk(cachedir):
+            for name in files:
+                if not name.endswith(cacheblock_suffix):
+                    continue
+
+                blockpath = os.path.join(root, name)
+                res = os.stat(blockpath)
+                usage += res.st_size
+        return usage
+
+
     @staticmethod
     def init_cache(cachedir, maxslots):
+        #
+        # First check the disk cache works at all by creating a 1 byte cache entry
+        #
+        checkexists = DiskCacheSlot.get_from_disk('0cc175b9c0f1b6a831c399e269772661', cachedir)
+        ds = DiskCacheSlot('0cc175b9c0f1b6a831c399e269772661', cachedir)
+        ds.set(b'a')
+        if checkexists is None:
+            # Don't keep the test entry around unless it existed beforehand.
+            ds.evict()
+
         # map in all the files in the cache directory, up to max slots.
         # after max slots, try to delete the excess blocks.
         #
@@ -208,13 +244,13 @@ class DiskCacheSlot(object):
 
         # Map in all the files we found, up to maxslots, if we exceed
         # maxslots, start throwing things out.
-        cachelist = []
+        cachelist: collections.OrderedDict = collections.OrderedDict()
         for b in blocks:
             got = DiskCacheSlot.get_from_disk(b[0], cachedir)
             if got is None:
                 continue
             if len(cachelist) < maxslots:
-                cachelist.append(got)
+                cachelist[got.locator] = got
             else:
                 # we found more blocks than maxslots, try to
                 # throw it out of the cache.