X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/563325b7d7336e3bfe9248b40c766d83e8b19bf3..10801e63c5029a2fffc3599235dddab62a2f3480:/sdk/python/arvados/diskcache.py

diff --git a/sdk/python/arvados/diskcache.py b/sdk/python/arvados/diskcache.py
index 9f218e56ce..f8fca57803 100644
--- a/sdk/python/arvados/diskcache.py
+++ b/sdk/python/arvados/diskcache.py
@@ -8,62 +8,88 @@ import os
 import traceback
 import stat
 import tempfile
-import hashlib
 import fcntl
+import time
+import errno
+import logging
+import weakref
+
+_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()
         return self.content
 
     def set(self, value):
+        tmpfile = None
         try:
             if value is None:
                 self.content = None
+                self.ready.set()
                 return
 
             if len(value) == 0:
                 # Can't mmap a 0 length file
                 self.content = b''
+                self.ready.set()
                 return
 
             if self.content is not None:
                 # Has been set already
+                self.ready.set()
                 return
 
             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)
+            final = os.path.join(blockdir, self.locator) + cacheblock_suffix
 
-            f = tempfile.NamedTemporaryFile(dir=blockdir, delete=False)
-            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 Exception as e:
-            traceback.print_exc()
-        finally:
+            self.content = mmap.mmap(self.filehandle.fileno(), 0, access=mmap.ACCESS_READ)
+            # only set the event when mmap is successful
             self.ready.set()
+        finally:
+            if tmpfile is not None:
+                # If the tempfile hasn't been renamed on disk yet, try to delete it.
+                try:
+                    os.remove(tmpfile)
+                except:
+                    pass
 
     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)
@@ -90,37 +116,63 @@ class DiskCacheSlot(object):
             # gone.
 
             blockdir = os.path.join(self.cachedir, self.locator[0:3])
-            final = os.path.join(blockdir, self.locator)
+            final = os.path.join(blockdir, self.locator) + cacheblock_suffix
             try:
-                # If we can't upgrade our shared lock to an exclusive
-                # lock, it'll throw an error, that's fine and
-                # desirable, it means another process has a lock and
-                # we shouldn't delete the block.
-                fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                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.
+                #
+                # 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(self.filehandle, fcntl.LOCK_EX | fcntl.LOCK_NB)
+
                 os.remove(final)
+                return True
             except OSError:
                 pass
+            finally:
+                self.filehandle = None
+                self.linger = weakref.ref(self.content)
+                self.content = None
+        return False
+
+    def gone(self):
+        # Test if an evicted object is lingering
+        return self.content is None and (self.linger is None or self.linger() is None)
 
     @staticmethod
     def get_from_disk(locator, cachedir):
-        # Get it, check it, return it
         blockdir = os.path.join(cachedir, locator[0:3])
-        final = os.path.join(blockdir, locator)
+        final = os.path.join(blockdir, locator) + cacheblock_suffix
 
         try:
             filehandle = open(final, "rb")
 
             # 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(filehandle, fcntl.LOCK_SH)
 
             content = mmap.mmap(filehandle.fileno(), 0, access=mmap.ACCESS_READ)
-            disk_md5 = hashlib.md5(content).hexdigest()
-            if disk_md5 == locator:
-                dc = DiskCacheSlot(locator, cachedir)
-                dc.content = content
-                dc.ready.set()
-                return dc
+            dc = DiskCacheSlot(locator, cachedir)
+            dc.filehandle = filehandle
+            dc.content = content
+            dc.ready.set()
+            return dc
         except FileNotFoundError:
             pass
         except Exception as e:
@@ -129,36 +181,72 @@ class DiskCacheSlot(object):
         return None
 
     @staticmethod
-    def cleanup_cachedir(cachedir, maxsize):
-        blocks = []
-        totalsize = 0
+    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)
-                blocks.append((blockpath, res.st_size, res.st_atime))
-                totalsize += res.st_size
+                usage += res.st_size
+        return usage
 
-        if totalsize <= maxsize:
-            return
 
-        # sort by atime, so the blocks accessed the longest time in
-        # the past get deleted first.
-        blocks.sort(key=lambda x: x[2])
+    @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.
+        #
+        # this gives the calling process ownership of all the blocks
 
-        # go through the list and try deleting blocks until we're
-        # below the target size and/or we run out of blocks
-        i = 0
-        while i < len(blocks) and totalsize > maxsize:
-            try:
-                with open(blocks[i][0], "rb") as f:
-                    # If we can't get an exclusive lock, it'll
-                    # throw an error, that's fine and desirable,
-                    # it means another process has a lock and we
-                    # shouldn't delete the block.
-                    fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
-                    os.remove(block)
-                    totalsize -= blocks[i][1]
-            except OSError:
-                pass
-            i += 1
+        blocks = []
+        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)
+
+                if len(name) == (32+len(cacheblock_suffix)) and not name.startswith("tmp"):
+                    blocks.append((name[0:32], res.st_atime))
+                elif name.startswith("tmp") and ((time.time() - res.st_mtime) > 60):
+                    # found a temporary file more than 1 minute old,
+                    # try to delete it.
+                    try:
+                        os.remove(blockpath)
+                    except:
+                        pass
+
+        # sort by access time (atime), going from most recently
+        # accessed (highest timestamp) to least recently accessed
+        # (lowest timestamp).
+        blocks.sort(key=lambda x: x[1], reverse=True)
+
+        # Map in all the files we found, up to maxslots, if we exceed
+        # maxslots, start throwing things out.
+        cachelist = []
+        for b in blocks:
+            got = DiskCacheSlot.get_from_disk(b[0], cachedir)
+            if got is None:
+                continue
+            if len(cachelist) < maxslots:
+                cachelist.append(got)
+            else:
+                # we found more blocks than maxslots, try to
+                # throw it out of the cache.
+                got.evict()
+
+        return cachelist