Share management of cache dir across multiple aptdir invocations. refs:

#52

moncic/deb.py

@@ -3,7 +3,7 @@
 import contextlib
 import os
 import tempfile
-from typing import Dict, Generator, NamedTuple
+from typing import Dict, Generator, NamedTuple, Optional
 
 
 class FileInfo(NamedTuple):
@@ -28,43 +28,59 @@ def __init__(self, cache_dir: str, cache_size: int = 512*1024*1024):
         self.cache_dir = cache_dir
         # Maximum cache size in bytes
         self.cache_size = cache_size
+        # Information about .deb files present in cache
+        self.debs: Dict[str, FileInfo] = {}
+        self.src_dir_fd: Optional[int] = None
+
+    def __enter__(self):
+        os.makedirs(self.cache_dir, exist_ok=True)
+        self.src_dir_fd = os.open(self.cache_dir, os.O_RDONLY)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # Do cache cleanup
+        self.trim_cache()
+
+        os.close(self.src_dir_fd)
+
+    def trim_cache(self):
+        """
+        Trim cache to fit self.cache_size, removing the files least recently
+        accessed
+        """
+        # Sort debs by atime and remove all those that go beyond
+        # self.cache_size
+        sdebs = sorted(self.debs.items(), key=lambda x: x[1].atime_ns, reverse=True)
+        size = 0
+        for name, info in sdebs:
+            if size > self.cache_size:
+                os.unlink(name, dir_fd=self.src_dir_fd)
+                del self.debs[name]
+            else:
+                size += info.size
 
     @contextlib.contextmanager
     def apt_archives(self) -> Generator[str, None, None]:
         """
         Create a directory that can be bind mounted as /apt/cache/apt/archives
         """
-        os.makedirs(self.cache_dir, exist_ok=True)
-        with dirfd(self.cache_dir) as src_dir_fd:
-            with tempfile.TemporaryDirectory(dir=self.cache_dir) as aptdir:
-                with dirfd(aptdir) as dst_dir_fd:
-                    debs: Dict[str, FileInfo] = {}
-
-                    # Handlink debs to temp dir
-                    with os.scandir(src_dir_fd) as it:
-                        for de in it:
-                            if de.name.endswith(".deb"):
-                                st = de.stat()
-                                debs[de.name] = FileInfo(st.st_size, st.st_atime_ns)
-                                os.link(de.name, de.name, src_dir_fd=src_dir_fd, dst_dir_fd=dst_dir_fd)
-
-                    yield aptdir
+        with tempfile.TemporaryDirectory(dir=self.cache_dir) as aptdir:
+            with dirfd(aptdir) as dst_dir_fd:
+                # Handlink debs to temp dir
+                with os.scandir(self.src_dir_fd) as it:
+                    for de in it:
+                        if de.name.endswith(".deb"):
+                            st = de.stat()
+                            self.debs[de.name] = FileInfo(st.st_size, st.st_atime_ns)
+                            os.link(de.name, de.name, src_dir_fd=self.src_dir_fd, dst_dir_fd=dst_dir_fd)
 
-                    # Hardlink new debs to cache dir
-                    with os.scandir(dst_dir_fd) as it:
-                        for de in it:
-                            if de.name.endswith(".deb"):
-                                st = de.stat()
-                                if de.name not in debs:
-                                    os.link(de.name, de.name, src_dir_fd=dst_dir_fd, dst_dir_fd=src_dir_fd)
-                                debs[de.name] = FileInfo(st.st_size, st.st_atime_ns)
+                yield aptdir
 
-                    # Sort debs by atime and remove all those that go beyond
-                    # self.cache_size
-                    sdebs = sorted(debs.items(), key=lambda x: x[1].atime_ns, reverse=True)
-                    size = 0
-                    for name, info in sdebs:
-                        if size > self.cache_size:
-                            os.unlink(name, dir_fd=src_dir_fd)
-                        else:
-                            size += info.size
+                # Hardlink new debs to cache dir
+                with os.scandir(dst_dir_fd) as it:
+                    for de in it:
+                        if de.name.endswith(".deb"):
+                            st = de.stat()
+                            if de.name not in self.debs:
+                                os.link(de.name, de.name, src_dir_fd=dst_dir_fd, dst_dir_fd=self.src_dir_fd)
+                            self.debs[de.name] = FileInfo(st.st_size, st.st_atime_ns)

tests/test_debs.py

@@ -19,11 +19,23 @@ def test_share(self):
         with tempfile.TemporaryDirectory() as workdir:
             make_deb(workdir, "a", 1000, 1)
             make_deb(workdir, "b", 2000, 2)
-            cache = DebCache(workdir, 5000)
-            with cache.apt_archives() as aptdir:
-                self.assertTrue(os.path.exists(os.path.join(aptdir, "a.deb")))
-                self.assertTrue(os.path.exists(os.path.join(aptdir, "b.deb")))
-                make_deb(aptdir, "c", 1000, 3)
-            self.assertTrue(os.path.exists(os.path.join(workdir, "a.deb")))
-            self.assertTrue(os.path.exists(os.path.join(workdir, "b.deb")))
-            self.assertTrue(os.path.exists(os.path.join(workdir, "c.deb")))
+            with DebCache(workdir, 5000) as cache:
+                with cache.apt_archives() as aptdir:
+                    self.assertTrue(os.path.exists(os.path.join(aptdir, "a.deb")))
+                    self.assertTrue(os.path.exists(os.path.join(aptdir, "b.deb")))
+                    make_deb(aptdir, "c", 1000, 3)
+                self.assertTrue(os.path.exists(os.path.join(workdir, "a.deb")))
+                self.assertTrue(os.path.exists(os.path.join(workdir, "b.deb")))
+                self.assertTrue(os.path.exists(os.path.join(workdir, "c.deb")))
+
+    def test_cache_limit(self):
+        with tempfile.TemporaryDirectory() as workdir:
+            make_deb(workdir, "a", 1000, 1)
+            make_deb(workdir, "b", 2000, 2)
+            with DebCache(workdir, 4000) as cache:
+                with cache.apt_archives() as aptdir:
+                    self.assertTrue(os.path.exists(os.path.join(aptdir, "a.deb")))
+                    self.assertTrue(os.path.exists(os.path.join(aptdir, "b.deb")))
+                    make_deb(aptdir, "c", 1500, 3)
+                self.assertTrue(os.path.exists(os.path.join(workdir, "b.deb")))
+                self.assertTrue(os.path.exists(os.path.join(workdir, "c.deb")))