1"""Bounded, authenticated disk caches with SSD promotion and request pinning."""
2from collections import OrderedDict
3import fcntl
4import logging
5import mmap
6import os
7from pathlib import Path
8import struct
9import threading
10
11from cryptography.exceptions import InvalidTag
12from cryptography.hazmat.primitives.ciphers.aead import AESGCM
13
14
15class DiskCache:
16 def __init__(self, ssd, hdd, ssd_bytes, hdd_bytes, key):
17 if min(ssd_bytes, hdd_bytes) <= 0:
18 raise ValueError("Disk cache capacities must be positive")
19 self.cipher = AESGCM(key)
20 self.roots = (Path(ssd), Path(hdd))
21 if (self.roots[0].resolve().is_relative_to(self.roots[1].resolve())
22 or self.roots[1].resolve().is_relative_to(self.roots[0].resolve())):
23 raise ValueError("Disk cache directories must be distinct")
24 self.limits = (ssd_bytes, hdd_bytes)
25 self.lock = threading.Lock()
26 self.writers = (threading.Lock(), threading.Lock())
27 self.entries = (OrderedDict(), OrderedDict())
28 self.pins = {}
29 self.files = []
30 try:
31 for tier, root in enumerate(self.roots):
32 root.mkdir(mode=0o700, parents=True, exist_ok=True)
33 root.chmod(0o700)
34 descriptor = os.open(root / ".lock", os.O_CREAT | os.O_RDWR, 0o600)
35 self.files.append(descriptor)
36 fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
37 for path in root.rglob("*.pending"):
38 path.unlink()
39 inventory = sorted((path.stat().st_mtime_ns, path) for path in root.rglob("*.bin"))
40 for _, path in inventory:
41 self.entries[tier][str(path.relative_to(root))] = path.stat().st_size
42 self._evict(tier, 0)
43 except Exception:
44 self.close()
45 raise
46
47 def close(self):
48 for descriptor in self.files:
49 os.close(descriptor)
50 self.files.clear()
51
52 def acquire(self, name, request):
53 with self.lock:
54 for entries in self.entries:
55 if name in entries:
56 entries.move_to_end(name)
57 self.pins.setdefault(name, set()).add(request)
58 return True
59 return False
60
61 def release(self, request):
62 with self.lock:
63 for name, requests in list(self.pins.items()):
64 requests.discard(request)
65 if not requests:
66 del self.pins[name]
67
68 def touch(self, name):
69 with self.lock:
70 for entries in self.entries:
71 if name in entries:
72 entries.move_to_end(name)
73 break
74
75 def _evict(self, tier, needed):
76 removed = []
77 with self.lock:
78 entries = self.entries[tier]
79 size = sum(entries.values())
80 for name in list(entries):
81 if size + needed <= self.limits[tier]:
82 break
83 if name not in self.pins:
84 size -= entries.pop(name)
85 removed.append(self.roots[tier] / name)
86 for path in removed:
87 self._unlink(tier, path)
88 return size + needed <= self.limits[tier]
89
90 def _unlink(self, tier, path):
91 path.unlink(missing_ok=True)
92 parent = path.parent
93 while parent != self.roots[tier]:
94 try:
95 parent.rmdir()
96 except OSError:
97 break
98 parent = parent.parent
99
100 def _put(self, tier, name, buffer):
101 path = self.roots[tier] / name
102 with self.writers[tier]:
103 with self.lock:
104 if name in self.entries[tier]:
105 self.entries[tier].move_to_end(name)
106 present = True
107 else:
108 present = False
109 if present:
110 os.utime(path, None)
111 return
112 if len(buffer) > self.limits[tier] or not self._evict(tier, len(buffer)):
113 return
114 path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
115 pending = path.with_suffix(".pending")
116 try:
117 descriptor = os.open(pending, os.O_CREAT | os.O_EXCL | os.O_WRONLY | getattr(os, "O_DIRECT", 0), 0o600)
118 try:
119 if os.write(descriptor, buffer) != len(buffer):
120 raise OSError("Short cache write")
121 os.fsync(descriptor)
122 finally:
123 os.close(descriptor)
124 pending.replace(path)
125 with self.lock:
126 self.entries[tier][name] = len(buffer)
127 finally:
128 pending.unlink(missing_ok=True)
129
130 def store(self, name, data):
131 # Page-aligned ciphertext keeps direct I/O on both ext4 and ZFS.
132 size = (len(data) + 36 + mmap.PAGESIZE - 1) // mmap.PAGESIZE * mmap.PAGESIZE
133 with mmap.mmap(-1, size) as buffer:
134 nonce = os.urandom(12)
135 buffer[:20] = struct.pack("!Q", len(data)) + nonce
136 view = memoryview(buffer)[20:20 + len(data) + 16]
137 try:
138 self.cipher.encrypt_into(nonce, data, name.encode(), view)
139 finally:
140 view.release()
141 self._put(1, name, buffer)
142 self._put(0, name, buffer)
143
144 def load(self, name, target):
145 pin = object()
146 if not self.acquire(name, pin):
147 raise FileNotFoundError("Cache block was evicted")
148 try:
149 for tier, root in enumerate(self.roots):
150 path = root / name
151 with self.lock:
152 size = self.entries[tier].get(name)
153 if size is None:
154 continue
155 try:
156 with mmap.mmap(-1, size) as buffer:
157 descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECT", 0))
158 try:
159 if os.readv(descriptor, [buffer]) != size:
160 raise OSError("Short cache read")
161 finally:
162 os.close(descriptor)
163 if struct.unpack("!Q", buffer[:8])[0] != len(target):
164 raise ValueError("Cache block size mismatch")
165 view = memoryview(buffer)[20:20 + len(target) + 16]
166 try:
167 self.cipher.decrypt_into(buffer[8:20], view, name.encode(), target)
168 finally:
169 view.release()
170 with self.lock:
171 self.entries[tier].move_to_end(name)
172 os.utime(path, None)
173 if tier == 1:
174 try:
175 self._put(0, name, buffer)
176 except OSError:
177 logging.getLogger(__name__).warning("SSD cache promotion failed", exc_info=True)
178 return "ssd" if tier == 0 else "hdd"
179 except (OSError, ValueError, struct.error, InvalidTag):
180 with self.writers[tier]:
181 with self.lock:
182 self.entries[tier].pop(name, None)
183 self._unlink(tier, path)
184 raise FileNotFoundError("No authenticated cache block remains")
185 finally:
186 self.release(pin)