diff --git a/src/mccode_antlr/reader/registry.py b/src/mccode_antlr/reader/registry.py index 278e4c7..cd4d675 100644 --- a/src/mccode_antlr/reader/registry.py +++ b/src/mccode_antlr/reader/registry.py @@ -68,7 +68,15 @@ def pooch_fetch(pooch_instance, filename: str) -> Path: directory = getattr(pooch_instance, 'abspath', None) if directory is not None: Path(directory, filename).parent.mkdir(parents=True, exist_ok=True) - return Path(pooch_instance.fetch(filename)) + for attempt in range(5): + try: + return Path(pooch_instance.fetch(filename)) + except PermissionError: + # On Windows, another process may be moving the same file into place + # (not fixed by the pooch PR above): + if attempt == 4: + raise + sleep(0.1 * (attempt + 1)) def _dedupe_identical_registry_entries(pooch_instance, names: list[str]) -> str | None: @@ -433,8 +441,15 @@ def _build_pooch(self): cache_path = pooch.os_cache(f'mccodeantlr/{safe_name}') registry_file_path = cache_path.joinpath(safe_version, safe_file) if registry_file_path.exists() and registry_file_path.is_file() and access(registry_file_path, R_OK): - with registry_file_path.open('r', encoding='utf-8') as file: - registry = {k: v for k, v in [x.strip().split(maxsplit=1) for x in file.readlines() if len(x)]} + for attempt in range(5): + try: + with registry_file_path.open('r', encoding='utf-8') as file: + registry = {k: v for k, v in [x.strip().split(maxsplit=1) for x in file.readlines() if len(x)]} + break + except PermissionError: # On Windows, while another process replaces it + if attempt == 4: + raise + sleep(0.1 * (attempt + 1)) else: # We allow a full-dictionary to be provided, otherwise we expect the registry file to be available from the # base_url where all subsequent files are also expected to be available @@ -450,9 +465,15 @@ def _build_pooch(self): last, check = check, check.parent # check is now a directory that exists, it may be the root of the filesystem if access(check, W_OK): + from os import getpid registry_file_path.parent.mkdir(parents=True, exist_ok=True) - with registry_file_path.open('w', encoding='utf-8') as file: - file.writelines('\n'.join([f'{k} {v}' for k, v in registry.items()])) + # Write-then-rename so a concurrent reader never sees a partial file. + tmp = registry_file_path.with_name(f'{registry_file_path.name}.{getpid()}.tmp') + tmp.write_text('\n'.join([f'{k} {v}' for k, v in registry.items()]), encoding='utf-8') + try: + tmp.replace(registry_file_path) + except OSError: # e.g. on Windows, if another process is reading it + tmp.unlink(missing_ok=True) else: logger.warning(f'Can not output {registry_file_path}, you lack write permissions for {check}') diff --git a/tests/test_pooch_fetch_parallel.py b/tests/test_pooch_fetch_parallel.py index 5304283..ef91524 100644 --- a/tests/test_pooch_fetch_parallel.py +++ b/tests/test_pooch_fetch_parallel.py @@ -48,3 +48,19 @@ def test_parallel_fetch_into_new_directories(tmp_path): assert [f for failed in results for f in failed] == [] for name in NAMES: assert (tmp_path / 'cache' / name).read_text() == name + + +def test_fetch_retried_on_permission_error(tmp_path): + class FlakyPooch: + abspath = tmp_path + calls = 0 + + def fetch(self, filename): + self.calls += 1 + if self.calls < 3: + raise PermissionError('in use by another process') + return str(tmp_path / filename) + + flaky = FlakyPooch() + assert pooch_fetch(flaky, 'd/f.txt') == tmp_path / 'd' / 'f.txt' + assert flaky.calls == 3 diff --git a/tests/test_registry_file_parallel.py b/tests/test_registry_file_parallel.py new file mode 100644 index 0000000..fafd97d --- /dev/null +++ b/tests/test_registry_file_parallel.py @@ -0,0 +1,22 @@ +"""Registries cached by several processes at the same time.""" +import multiprocessing +from pathlib import Path + +N = 20000 + + +def _registry_size(cache): + import pooch + from mccode_antlr.reader.registry import GitHubRegistry + pooch.os_cache = lambda name: Path(cache, name) + files = {f'dir{i}/Comp{i}.comp': 'sha256:' + 'ab' * 32 for i in range(N)} + try: + return len(GitHubRegistry('test', 'https://example.invalid/x', 'v1.0', registry=files).pooch.registry) + except Exception as error: + return repr(error) + + +def test_parallel_registry_file(tmp_path): + with multiprocessing.get_context('spawn').Pool(8) as pool: + for i in range(5): + assert pool.map(_registry_size, [str(tmp_path / str(i))] * 8) == [N] * 8