Merge branch 'lrucache' of https://github.com/sourcejedi/borg

borg/lrucache.py

@@ -1,42 +1,41 @@
-class LRUCache(dict):
-
-    def __init__(self, capacity):
-        super().__init__()
+class LRUCache:
+    def __init__(self, capacity, dispose):
+        self._cache = {}
         self._lru = []
         self._capacity = capacity
+        self._dispose = dispose
 
     def __setitem__(self, key, value):
-        try:
-            self._lru.remove(key)
-        except ValueError:
-            pass
+        assert key not in self._cache, (
+            "Unexpected attempt to replace a cached item,"
+            " without first deleting the old item.")
         self._lru.append(key)
         while len(self._lru) > self._capacity:
             del self[self._lru[0]]
-        return super().__setitem__(key, value)
+        self._cache[key] = value
 
     def __getitem__(self, key):
-        try:
-            self._lru.remove(key)
-            self._lru.append(key)
-        except ValueError:
-            pass
-        return super().__getitem__(key)
+        value = self._cache[key]  # raise KeyError if not found
+        self._lru.remove(key)
+        self._lru.append(key)
+        return value
 
     def __delitem__(self, key):
-        try:
-            self._lru.remove(key)
-        except ValueError:
-            pass
-        return super().__delitem__(key)
+        value = self._cache.pop(key)  # raise KeyError if not found
+        self._dispose(value)
+        self._lru.remove(key)
 
-    def pop(self, key, default=None):
-        try:
-            self._lru.remove(key)
-        except ValueError:
-            pass
-        return super().pop(key, default)
+    def __contains__(self, key):
+        return key in self._cache
 
-    def _not_implemented(self, *args, **kw):
-        raise NotImplementedError
-    popitem = setdefault = update = _not_implemented
+    def clear(self):
+        for value in self._cache.values():
+            self._dispose(value)
+        self._cache.clear()
+
+    # useful for testing
+    def items(self):
+        return self._cache.items()
+
+    def __len__(self):
+        return len(self._cache)

borg/repository.py

@@ -432,7 +432,8 @@ class LoggedIO:
 
     def __init__(self, path, limit, segments_per_dir, capacity=90):
         self.path = path
-        self.fds = LRUCache(capacity)
+        self.fds = LRUCache(capacity,
+                            dispose=lambda fd: fd.close())
         self.segment = 0
         self.limit = limit
         self.segments_per_dir = segments_per_dir
@@ -440,9 +441,8 @@ def __init__(self, path, limit, segments_per_dir, capacity=90):
         self._write_fd = None
 
     def close(self):
-        for segment in list(self.fds.keys()):
-            self.fds.pop(segment).close()
         self.close_segment()
+        self.fds.clear()
         self.fds = None  # Just to make sure we're disabled
 
     def segment_iterator(self, reverse=False):
@@ -516,9 +516,8 @@ def get_fd(self, segment):
             return fd
 
     def delete_segment(self, segment):
-        fd = self.fds.pop(segment)
-        if fd is not None:
-            fd.close()
+        if segment in self.fds:
+            del self.fds[segment]
         try:
             os.unlink(self.segment_filename(segment))
         except OSError:
@@ -561,9 +560,8 @@ def iter_objects(self, segment, include_data=False):
             header = fd.read(self.header_fmt.size)
 
     def recover_segment(self, segment, filename):
-        fd = self.fds.pop(segment)
-        if fd is not None:
-            fd.close()
+        if segment in self.fds:
+            del self.fds[segment]
         # FIXME: save a copy of the original file
         with open(filename, 'rb') as fd:
             data = memoryview(fd.read())

borg/testsuite/lrucache.py

@@ -1,40 +1,52 @@
 from ..lrucache import LRUCache
-from . import BaseTestCase
+import pytest
+from tempfile import TemporaryFile
 
 
-class LRUCacheTestCase(BaseTestCase):
+class TestLRUCache:
 
-    def test(self):
-        c = LRUCache(2)
-        self.assert_equal(len(c), 0)
+    def test_lrucache(self):
+        c = LRUCache(2, dispose=lambda _: None)
+        assert len(c) == 0
+        assert c.items() == set()
         for i, x in enumerate('abc'):
             c[x] = i
-        self.assert_equal(len(c), 2)
-        self.assert_equal(set(c), set(['b', 'c']))
-        self.assert_equal(set(c.items()), set([('b', 1), ('c', 2)]))
-        self.assert_equal(False, 'a' in c)
-        self.assert_equal(True, 'b' in c)
-        self.assert_raises(KeyError, lambda: c['a'])
-        self.assert_equal(c['b'], 1)
-        self.assert_equal(c['c'], 2)
+        assert len(c) == 2
+        assert c.items() == set([('b', 1), ('c', 2)])
+        assert 'a' not in c
+        assert 'b' in c
+        with pytest.raises(KeyError):
+            c['a']
+        assert c['b'] == 1
+        assert c['c'] == 2
         c['d'] = 3
-        self.assert_equal(len(c), 2)
-        self.assert_equal(c['c'], 2)
-        self.assert_equal(c['d'], 3)
-        c['c'] = 22
-        c['e'] = 4
-        self.assert_equal(len(c), 2)
-        self.assert_raises(KeyError, lambda: c['d'])
-        self.assert_equal(c['c'], 22)
-        self.assert_equal(c['e'], 4)
+        assert len(c) == 2
+        assert c['c'] == 2
+        assert c['d'] == 3
         del c['c']
-        self.assert_equal(len(c), 1)
-        self.assert_raises(KeyError, lambda: c['c'])
-        self.assert_equal(c['e'], 4)
+        assert len(c) == 1
+        with pytest.raises(KeyError):
+            c['c']
+        assert c['d'] == 3
+        c.clear()
+        assert c.items() == set()
 
-    def test_pop(self):
-        c = LRUCache(2)
-        c[1] = 1
-        c[2] = 2
-        c.pop(1)
-        c[3] = 3
+    def test_dispose(self):
+        c = LRUCache(2, dispose=lambda f: f.close())
+        f1 = TemporaryFile()
+        f2 = TemporaryFile()
+        f3 = TemporaryFile()
+        c[1] = f1
+        c[2] = f2
+        assert not f2.closed
+        c[3] = f3
+        assert 1 not in c
+        assert f1.closed
+        assert 2 in c
+        assert not f2.closed
+        del c[2]
+        assert 2 not in c
+        assert f2.closed
+        c.clear()
+        assert not c.items()
+        assert f3.closed