arvados.diskcache

  1# Copyright (C) The Arvados Authors. All rights reserved.
  2#
  3# SPDX-License-Identifier: Apache-2.0
  4
  5import threading
  6import mmap
  7import os
  8import traceback
  9import stat
 10import tempfile
 11import fcntl
 12import time
 13import errno
 14import logging
 15import weakref
 16import collections
 17
 18_logger = logging.getLogger('arvados.keep')
 19
 20cacheblock_suffix = ".keepcacheblock"
 21
 22class DiskCacheSlot(object):
 23    __slots__ = ("locator", "ready", "content", "cachedir", "filehandle", "linger")
 24
 25    def __init__(self, locator, cachedir):
 26        self.locator = locator
 27        self.ready = threading.Event()
 28        self.content = None
 29        self.cachedir = cachedir
 30        self.filehandle = None
 31        self.linger = None
 32
 33    def get(self):
 34        self.ready.wait()
 35        # 'content' can None, an empty byte string, or a nonempty mmap
 36        # region.  If it is an mmap region, we want to advise the
 37        # kernel we're going to use it.  This nudges the kernel to
 38        # re-read most or all of the block if necessary (instead of
 39        # just a few pages at a time), reducing the number of page
 40        # faults and improving performance by 4x compared to not
 41        # calling madvise.
 42        if self.content and hasattr(mmap.mmap, 'madvise'):
 43            self.content.madvise(mmap.MADV_WILLNEED)
 44        return self.content
 45
 46    def set(self, value):
 47        tmpfile = None
 48        try:
 49            if value is None:
 50                self.content = None
 51                self.ready.set()
 52                return False
 53
 54            if len(value) == 0:
 55                # Can't mmap a 0 length file
 56                self.content = b''
 57                self.ready.set()
 58                return True
 59
 60            if self.content is not None:
 61                # Has been set already
 62                self.ready.set()
 63                return False
 64
 65            blockdir = os.path.join(self.cachedir, self.locator[0:3])
 66            os.makedirs(blockdir, mode=0o700, exist_ok=True)
 67
 68            final = os.path.join(blockdir, self.locator) + cacheblock_suffix
 69
 70            self.filehandle = tempfile.NamedTemporaryFile(dir=blockdir, delete=False, prefix="tmp", suffix=cacheblock_suffix)
 71            tmpfile = self.filehandle.name
 72            os.chmod(tmpfile, stat.S_IRUSR | stat.S_IWUSR)
 73
 74            # aquire a shared lock, this tells other processes that
 75            # we're using this block and to please not delete it.
 76            fcntl.flock(self.filehandle, fcntl.LOCK_SH)
 77
 78            self.filehandle.write(value)
 79            self.filehandle.flush()
 80            os.rename(tmpfile, final)
 81            tmpfile = None
 82
 83            self.content = mmap.mmap(self.filehandle.fileno(), 0, access=mmap.ACCESS_READ)
 84            # only set the event when mmap is successful
 85            self.ready.set()
 86            return True
 87        finally:
 88            if tmpfile is not None:
 89                # If the tempfile hasn't been renamed on disk yet, try to delete it.
 90                try:
 91                    os.remove(tmpfile)
 92                except:
 93                    pass
 94
 95    def size(self):
 96        if self.content is None:
 97            if self.linger is not None:
 98                # If it is still lingering (object is still accessible
 99                # through the weak reference) it is still taking up
100                # space.
101                content = self.linger()
102                if content is not None:
103                    return len(content)
104            return 0
105        else:
106            return len(self.content)
107
108    def evict(self):
109        if not self.content:
110            return
111
112        # The mmap region might be in use when we decided to evict
113        # it.  This can happen if the cache is too small.
114        #
115        # If we call close() now, it'll throw an error if
116        # something tries to access it.
117        #
118        # However, we don't need to explicitly call mmap.close()
119        #
120        # I confirmed in mmapmodule.c that that both close
121        # and deallocate do the same thing:
122        #
123        # a) close the file descriptor
124        # b) unmap the memory range
125        #
126        # So we can forget it in the cache and delete the file on
127        # disk, and it will tear it down after any other
128        # lingering Python references to the mapped memory are
129        # gone.
130
131        blockdir = os.path.join(self.cachedir, self.locator[0:3])
132        final = os.path.join(blockdir, self.locator) + cacheblock_suffix
133        try:
134            fcntl.flock(self.filehandle, fcntl.LOCK_UN)
135
136            # try to get an exclusive lock, this ensures other
137            # processes are not using the block.  It is
138            # nonblocking and will throw an exception if we
139            # can't get it, which is fine because that means
140            # we just won't try to delete it.
141            #
142            # I should note here, the file locking is not
143            # strictly necessary, we could just remove it and
144            # the kernel would ensure that the underlying
145            # inode remains available as long as other
146            # processes still have the file open.  However, if
147            # you have multiple processes sharing the cache
148            # and deleting each other's files, you'll end up
149            # with a bunch of ghost files that don't show up
150            # in the file system but are still taking up
151            # space, which isn't particularly user friendly.
152            # The locking strategy ensures that cache blocks
153            # in use remain visible.
154            #
155            fcntl.flock(self.filehandle, fcntl.LOCK_EX | fcntl.LOCK_NB)
156
157            os.remove(final)
158            return True
159        except OSError:
160            pass
161        finally:
162            self.filehandle = None
163            self.content = None
164
165    @staticmethod
166    def get_from_disk(locator, cachedir):
167        blockdir = os.path.join(cachedir, locator[0:3])
168        final = os.path.join(blockdir, locator) + cacheblock_suffix
169
170        try:
171            filehandle = open(final, "rb")
172
173            # aquire a shared lock, this tells other processes that
174            # we're using this block and to please not delete it.
175            fcntl.flock(filehandle, fcntl.LOCK_SH)
176
177            content = mmap.mmap(filehandle.fileno(), 0, access=mmap.ACCESS_READ)
178            dc = DiskCacheSlot(locator, cachedir)
179            dc.filehandle = filehandle
180            dc.content = content
181            dc.ready.set()
182            return dc
183        except FileNotFoundError:
184            pass
185        except Exception as e:
186            traceback.print_exc()
187
188        return None
189
190    @staticmethod
191    def cache_usage(cachedir):
192        usage = 0
193        for root, dirs, files in os.walk(cachedir):
194            for name in files:
195                if not name.endswith(cacheblock_suffix):
196                    continue
197
198                blockpath = os.path.join(root, name)
199                res = os.stat(blockpath)
200                usage += res.st_size
201        return usage
202
203
204    @staticmethod
205    def init_cache(cachedir, maxslots):
206        #
207        # First check the disk cache works at all by creating a 1 byte cache entry
208        #
209        checkexists = DiskCacheSlot.get_from_disk('0cc175b9c0f1b6a831c399e269772661', cachedir)
210        ds = DiskCacheSlot('0cc175b9c0f1b6a831c399e269772661', cachedir)
211        ds.set(b'a')
212        if checkexists is None:
213            # Don't keep the test entry around unless it existed beforehand.
214            ds.evict()
215
216        # map in all the files in the cache directory, up to max slots.
217        # after max slots, try to delete the excess blocks.
218        #
219        # this gives the calling process ownership of all the blocks
220
221        blocks = []
222        for root, dirs, files in os.walk(cachedir):
223            for name in files:
224                if not name.endswith(cacheblock_suffix):
225                    continue
226
227                blockpath = os.path.join(root, name)
228                res = os.stat(blockpath)
229
230                if len(name) == (32+len(cacheblock_suffix)) and not name.startswith("tmp"):
231                    blocks.append((name[0:32], res.st_atime))
232                elif name.startswith("tmp") and ((time.time() - res.st_mtime) > 60):
233                    # found a temporary file more than 1 minute old,
234                    # try to delete it.
235                    try:
236                        os.remove(blockpath)
237                    except:
238                        pass
239
240        # sort by access time (atime), going from most recently
241        # accessed (highest timestamp) to least recently accessed
242        # (lowest timestamp).
243        blocks.sort(key=lambda x: x[1], reverse=True)
244
245        # Map in all the files we found, up to maxslots, if we exceed
246        # maxslots, start throwing things out.
247        cachelist: collections.OrderedDict = collections.OrderedDict()
248        for b in blocks:
249            got = DiskCacheSlot.get_from_disk(b[0], cachedir)
250            if got is None:
251                continue
252            if len(cachelist) < maxslots:
253                cachelist[got.locator] = got
254            else:
255                # we found more blocks than maxslots, try to
256                # throw it out of the cache.
257                got.evict()
258
259        return cachelist