Skip to content

Commit e15debd

Browse files
committed
Reuse a single per-file diskcache handle instead of reopening it per read block
ChunkedFile opened a fresh diskcache.Cache on every cache access, and while streaming a proxied file RemoteFile spawns a FileDownload per chunk, so serving one file constructed the cache O(file size) times. Each construction opens a SQLite connection (pragmas + fsync), negligible on SSD/NVMe but collapsing proxied browsing on SD-card devices: a 13.7MB file took ~8 minutes under concurrency and 34KB thumbnails ~50s; both drop to ~1-4s with this change. Memoize the Cache handle on the ChunkedFile and reuse it for the object's lifetime, and let a RemoteFile share its handle with the FileDownload streaming into the same file, so a full stream constructs the cache once regardless of size. To keep the memoized handle correct when the chunk directory is deleted and recreated (e.g. streamed cache eviction racing a stream), _open_cache stamps each handle with a per-incarnation token - the directory inode, or a marker file on filesystems without a usable inode (FAT/exFAT, some Android) - and reopens when it changes, restoring the self-healing that opening per access gave for free. Also stop resetting the read position in _initialize (now reachable mid-read via ensure_writable, where it would rewind an in-progress read), and skip rather than abort on a directory that cannot be removed during eviction (a held-open handle blocks removal on Windows). Claude-Session: https://claude.ai/code/session_01MnAmUbaCBpL4wBKFYUyCwt
1 parent fe61172 commit e15debd

1 file changed

Lines changed: 137 additions & 5 deletions

File tree

kolibri/utils/file_transfer.py

Lines changed: 137 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,11 @@ class ChunkedFileDoesNotExist(Exception):
101101

102102
CHUNK_SUFFIX = ".chunks"
103103

104+
# Marker file used to detect when a chunk directory has been deleted and
105+
# recreated, on filesystems that do not report a usable inode (see
106+
# ChunkedFile._chunk_dir_incarnation).
107+
INCARNATION_MARKER = ".incarnation"
108+
104109

105110
class TransferFileBase(BufferedIOBase, ABC):
106111
"""Abstract base class for file transfer destination objects."""
@@ -182,8 +187,17 @@ def _do_file_eviction(self, chunked_file_stats, file_size):
182187
if file_size <= evicted_file_size:
183188
break
184189
file_stats = chunked_file_stats[chunked_file_dir]
190+
try:
191+
shutil.rmtree(chunked_file_dir)
192+
except OSError as e:
193+
# The directory may be in use - for example a diskcache handle
194+
# held open while a file is being streamed, which blocks removal
195+
# on Windows. Skip it rather than aborting the whole eviction.
196+
logger.warning(
197+
"Could not evict chunked file {}: {}".format(chunked_file_dir, e)
198+
)
199+
continue
185200
evicted_file_size += file_stats["size"]
186-
shutil.rmtree(chunked_file_dir)
187201
return evicted_file_size
188202

189203
def evict_files(self, file_size):
@@ -213,9 +227,24 @@ class ChunkedFile(TransferFileBase):
213227
# Set chunk size to 128KB
214228
chunk_size = 128 * 1024
215229

216-
def __init__(self, filepath, raise_if_empty=False, raise_if_exists=False):
230+
def __init__(
231+
self, filepath, raise_if_empty=False, raise_if_exists=False, cache=None
232+
):
217233
self.filepath = filepath
218234
self.chunk_dir = filepath + CHUNK_SUFFIX
235+
# The diskcache handle is opened lazily and then reused for the lifetime
236+
# of this object, rather than being reopened for every cache access.
237+
# Opening a diskcache constructs a SQLite connection (with pragmas and
238+
# fsyncs), which is cheap on fast storage but very expensive on slow
239+
# storage like SD cards, where reopening it per read block made serving
240+
# proxied remote files pathologically slow. A caller may inject an
241+
# existing handle (see ``cache``) so that related objects for the same
242+
# file (e.g. a RemoteFile and the FileDownload streaming into it) share
243+
# one handle; an injected handle is not owned and so is not closed here.
244+
# These are assigned before the early returns below so that close() on a
245+
# partially constructed instance does not raise.
246+
self._cache = cache
247+
self._owns_cache = cache is None
219248
if raise_if_empty and not os.path.exists(self.chunk_dir):
220249
raise FileNotFoundError("Chunked file does not exist")
221250
if raise_if_exists and os.path.exists(self.filepath):
@@ -224,19 +253,96 @@ def __init__(self, filepath, raise_if_empty=False, raise_if_exists=False):
224253
self.position = 0
225254
self._file_size = None
226255

256+
def _chunk_dir_incarnation(self):
257+
"""
258+
Return a token identifying the current on-disk incarnation of the chunk
259+
directory, so that a deletion + recreation (e.g. streamed cache
260+
eviction) can be detected and a stale diskcache handle dropped. Prefers
261+
the directory inode (a cheap stat), and falls back to a marker file on
262+
filesystems that do not report a usable inode (e.g. FAT/exFAT, some
263+
Android storage), where st_ino is 0.
264+
"""
265+
try:
266+
inode = os.stat(self.chunk_dir).st_ino
267+
if inode:
268+
# The inode changes when the directory is deleted and recreated.
269+
# We deliberately do not pair it with mtime/ctime: those change
270+
# whenever a chunk file is written into the directory, which
271+
# would churn the token during normal streaming. The residual
272+
# risk is a recreated directory being assigned the exact same
273+
# inode number (POSIX inode recycling) while a stale handle is
274+
# still held - vanishingly unlikely, and only reachable in the
275+
# already-narrow eviction-mid-stream race.
276+
return inode
277+
marker = os.path.join(self.chunk_dir, INCARNATION_MARKER)
278+
try:
279+
with open(marker) as f:
280+
return f.read()
281+
except FileNotFoundError:
282+
token = os.urandom(8).hex()
283+
try:
284+
# Atomic create so concurrent openers agree on one token.
285+
with open(marker, "x") as f:
286+
f.write(token)
287+
return token
288+
except FileExistsError:
289+
with open(marker) as f:
290+
return f.read()
291+
except OSError:
292+
# The directory was removed between the _check_for_chunk_dir guard
293+
# and here (eviction racing a stream). Normalize to the exception
294+
# callers already handle, rather than letting a bare OSError escape
295+
# and fail the download hard.
296+
raise ChunkedFileDoesNotExist("Chunked file does not exist")
297+
298+
@contextmanager
227299
def _open_cache(self):
228300
self._check_for_chunk_dir()
229-
return Cache(self.cache_dir)
301+
incarnation = self._chunk_dir_incarnation()
302+
if self._cache is not None and (
303+
getattr(self._cache, "_kolibri_chunk_incarnation", None) != incarnation
304+
):
305+
# The chunk directory was deleted and recreated (e.g. by streamed
306+
# cache eviction) since this handle was opened, so it points at the
307+
# old, now-deleted diskcache. Drop it and reopen against the live
308+
# directory, matching the pre-memoization behaviour of always using
309+
# the current cache. This matters for correctness, not just leaks:
310+
# chunk locks live as rows in the on-disk SQLite DB, so a stale
311+
# handle would silently fail to coordinate with other readers.
312+
self._close_cache()
313+
if self._cache is None:
314+
cache = Cache(self.cache_dir)
315+
# Stamp the handle with the incarnation it was opened against so a
316+
# later deletion/recreation can be detected.
317+
cache._kolibri_chunk_incarnation = incarnation
318+
self._cache = cache
319+
# We opened this handle, so this instance owns and must close it.
320+
self._owns_cache = True
321+
yield self._cache
322+
323+
def get_cache(self):
324+
"""
325+
Return the reused diskcache handle for this file, opening it if needed.
326+
Used to share a single handle with a FileDownload for the same file.
327+
"""
328+
with self._open_cache() as cache:
329+
return cache
230330

231331
def _initialize(self):
232332
os.makedirs(self.chunk_dir, exist_ok=True)
233333
self.cache_dir = os.path.join(self.chunk_dir, ".cache")
234-
self.position = 0
334+
# NB: do not reset self.position here. _initialize is also invoked by
335+
# ensure_writable() to recreate a directory that was cleaned up mid-read,
336+
# and resetting the read cursor there would silently rewind an
337+
# in-progress read. __init__ sets the initial position separately.
235338

236339
def ensure_writable(self):
237340
try:
238341
self._check_for_chunk_dir()
239342
except ChunkedFileDoesNotExist:
343+
# The chunk directory was removed by another process; recreate it.
344+
# Any cache handle we still hold is detected as stale and reopened on
345+
# next use (see _open_cache), so we don't need to touch it here.
240346
self._initialize()
241347

242348
@property
@@ -474,8 +580,15 @@ def md5_checksum(self):
474580
def delete(self):
475581
shutil.rmtree(self.chunk_dir)
476582

583+
def _close_cache(self):
584+
# Only close a handle we opened ourselves; an injected handle is owned
585+
# (and closed) by whoever provided it.
586+
if self._owns_cache and self._cache is not None:
587+
self._cache.close()
588+
self._cache = None
589+
477590
def close(self):
478-
pass
591+
self._close_cache()
479592

480593
def __enter__(self):
481594
return self
@@ -772,13 +885,21 @@ def __init__(
772885
timeout=Transfer.DEFAULT_TIMEOUT,
773886
retry_wait=30,
774887
full_ranges=True,
888+
cache=None,
775889
):
776890

777891
# Allow an existing requests.Session to be passed in, so it can be
778892
# reused for speed. The default blocks cross-host redirects so a
779893
# caller-supplied baseurl can't pivot us onto another host.
780894
self.session = session or SameHostSession()
781895

896+
# Allow an existing diskcache handle to be passed in, so the chunked
897+
# destination file reuses it rather than opening its own. This lets a
898+
# RemoteFile streaming a proxied file share one handle across itself and
899+
# every FileDownload it spawns, instead of reopening the diskcache once
900+
# per downloaded chunk.
901+
self._cache = cache
902+
782903
# A flag to allow the download to remain in the chunked file directory
783904
# for easier clean up when it is just a temporary download.
784905
self._finalize_download = finalize_download
@@ -815,6 +936,7 @@ def _initialize_dest_file(self):
815936
self.dest,
816937
raise_if_empty=self.full_ranges,
817938
raise_if_exists=self.full_ranges,
939+
cache=self._cache,
818940
)
819941
except FileNotFoundError:
820942
# No chunked file exists, use TransferFile for direct download
@@ -1133,13 +1255,21 @@ def _run_transfer(self):
11331255

11341256
def _start_transfer(self, start=None, end=None):
11351257
if not self.is_complete(start=start, end=end):
1258+
# Recreate the chunk directory (and drop any stale cache handle) if
1259+
# it was cleaned up by another process since we opened it, so that
1260+
# sharing our cache handle below opens against a directory that
1261+
# exists rather than raising.
1262+
self.ensure_writable()
11361263
self.transfer = FileDownload(
11371264
self.remote_url,
11381265
self.filepath,
11391266
start_range=start,
11401267
end_range=end,
11411268
finalize_download=False,
11421269
full_ranges=False,
1270+
# Share this file's diskcache handle with the download, so
1271+
# streaming a file doesn't reopen the diskcache once per chunk.
1272+
cache=self.get_cache(),
11431273
)
11441274
with self._open_cache() as cache:
11451275
header_info = cache.get(self.remote_url)
@@ -1171,3 +1301,5 @@ def close(self):
11711301
self.transfer.close()
11721302
if self._dest_file_handle:
11731303
self._dest_file_handle.close()
1304+
# Close the reused diskcache handle shared with any FileDownload.
1305+
super().close()

0 commit comments

Comments
 (0)