From fbce0fc85fcb0a87998567163a1802fa3c8cf370 Mon Sep 17 00:00:00 2001 From: Aaron Kanzer Date: Sat, 29 Aug 2026 11:06:51 -0400 Subject: [PATCH] fix: coordinate concurrent first downloads of the same cache key resolve()'s check-then-download wasn't coordinated across threads: N threads racing the first resolve() of a key (e.g. a thread pool calling chembl_bridge() concurrently right after enable_cache()) each saw it missing and each downloaded it in full. Reproduced directly: 16 threads calling chembl_bridge() concurrently on an empty cache logged 16 separate "caching ..." downloads of the same file instead of one. Fixed with a per-key lock: the first thread to reach a key downloads it, the rest wait and then reuse the file it wrote. Verified the same repro now logs exactly one download. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01L4ztUhMgoMFRicB8uiYx76 --- README.md | 2 +- pyproject.toml | 2 +- src/scigantic_bindingdb/cache.py | 34 ++++++++++++++++++++---- tests/test_cache.py | 44 ++++++++++++++++++++++++++++++++ 4 files changed, 75 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 84eefd5..9c9770b 100644 --- a/README.md +++ b/README.md @@ -105,7 +105,7 @@ bindingdb.enable_cache() df = bindingdb.chembl_bridge() # downloads the bridge table once, then reads from disk ``` -`chembl_bridge()` and `dti_pairs()` each need exactly one derived file, so caching downloads that one file to `~/.cache/scigantic-bindingdb` (override with `enable_cache(cache_dir=...)` or the `SCIGANTIC_BINDINGDB_CACHE` environment variable) and reuses it after that. +`chembl_bridge()` and `dti_pairs()` each need exactly one derived file, so caching downloads that one file to `~/.cache/scigantic-bindingdb` (override with `enable_cache(cache_dir=...)` or the `SCIGANTIC_BINDINGDB_CACHE` environment variable) and reuses it after that. Concurrent callers racing the first download of the same file wait for it rather than each downloading their own copy. `connect()`, `query()` and `measurements()` don't participate in this: `connect()` registers five core tables as views the first time a release is used, so caching them there would mean the first call for any release eagerly downloads everything regardless of what the query actually touches. Cache one table yourself if you want it locally: `bindingdb.cache_resolve("202608/parquet/measurements.parquet")` downloads it and returns the local path, usable directly in `read_parquet(...)`. diff --git a/pyproject.toml b/pyproject.toml index 5593662..ef3a9a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "scigantic-bindingdb" -version = "0.2.3" +version = "0.2.4" description = "Query BindingDB directly from a public S3 mirror with DuckDB, including a ChEMBL cross-reference bridge table and a ready drug-target-interaction training table." readme = "README.md" requires-python = ">=3.10" diff --git a/src/scigantic_bindingdb/cache.py b/src/scigantic_bindingdb/cache.py index cf1dd91..95140b0 100644 --- a/src/scigantic_bindingdb/cache.py +++ b/src/scigantic_bindingdb/cache.py @@ -23,6 +23,7 @@ import os import sys +import threading import urllib.request import uuid from pathlib import Path @@ -34,6 +35,24 @@ _CHUNK_BYTES = 1024 * 1024 +# One lock per key, created on first use. Guards resolve()'s +# check-then-download against concurrent callers asking for the same key +# at once: without this, N threads racing the first resolve() of a key +# each see it missing and each download it in full, rather than one +# downloading while the rest wait and reuse the result. Verified directly: +# 16 threads calling chembl_bridge() concurrently on an empty cache +# triggered 16 separate downloads of the same file before this fix. +_resolve_locks: dict[str, threading.Lock] = {} +_resolve_locks_guard = threading.Lock() + + +def _lock_for(key: str) -> threading.Lock: + lock = _resolve_locks.get(key) + if lock is not None: + return lock + with _resolve_locks_guard: + return _resolve_locks.setdefault(key, threading.Lock()) + def _default_cache_dir() -> Path: if sys.platform == "win32": @@ -114,7 +133,9 @@ def resolve(key: str) -> str: `key` is a path relative to the bucket root, e.g. "202608/derived/bindingdb_chembl_bridge.parquet". Downloads to the cache on first access; later calls for the same key reuse the local - file without touching the network. + file without touching the network. Concurrent callers asking for the + same key while it's still downloading wait for that download rather + than each starting their own. """ if not _enabled: return f"s3://{BUCKET}/{key}" @@ -124,8 +145,11 @@ def resolve(key: str) -> str: if local_path.exists(): return str(local_path) - local_path.parent.mkdir(parents=True, exist_ok=True) - url = f"https://{BUCKET}.s3.{REGION}.amazonaws.com/{key}" - print(f"scigantic-bindingdb: caching {key} ...", file=sys.stderr, flush=True) - _atomic_download(url, local_path) + with _lock_for(key): + if local_path.exists(): # another thread finished it while we waited + return str(local_path) + local_path.parent.mkdir(parents=True, exist_ok=True) + url = f"https://{BUCKET}.s3.{REGION}.amazonaws.com/{key}" + print(f"scigantic-bindingdb: caching {key} ...", file=sys.stderr, flush=True) + _atomic_download(url, local_path) return str(local_path) diff --git a/tests/test_cache.py b/tests/test_cache.py index 7b14278..0466e4c 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -127,6 +127,50 @@ def download(_i): assert list(local_path.parent.glob("key.parquet.*.part")) == [] +def test_concurrent_first_resolve_of_the_same_key_downloads_once(tmp_path): + # Regression test for a real inefficiency: resolve()'s check-then- + # download wasn't coordinated across threads, so N threads racing the + # first resolve() of a key each saw it missing and each downloaded it + # in full. Verified directly before this was fixed with a per-key + # lock: 16 threads calling resolve() concurrently on an empty cache + # triggered 16 separate downloads of the same file instead of one. + # + # mock.patch is entered/exited only once here, by the main thread, + # before the pool starts; the worker threads only read the already- + # patched attribute, so this doesn't hit the same-target mock.patch + # race that test_concurrent_downloads_of_the_same_key_never_raise's + # docstring describes for patching from multiple threads at once. + body = b"cached-body" + + class FakeConnection: + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + def read(self, *_args): + if getattr(self, "_served", False): + return b"" + self._served = True + return body + + bindingdb.enable_cache(cache_dir=str(tmp_path)) + try: + with mock.patch( + "urllib.request.urlopen", side_effect=lambda *a, **k: FakeConnection() + ) as urlopen: + with ThreadPoolExecutor(max_workers=16) as pool: + paths = list(pool.map(lambda _i: resolve("shared/once.parquet"), range(16))) + + assert len(set(paths)) == 1 + assert urlopen.call_count == 1 # only one thread actually downloaded + final = tmp_path / "shared" / "once.parquet" + assert final.read_bytes() == body + finally: + bindingdb.disable_cache() + + def test_cached_bridge_matches_uncached(tmp_path): # The one integration-level check: caching actually gets used by a real # function, not just by resolve() directly, and returns the same data.