Write the cached registry file atomically

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}')
 

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

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