2/3 repair_wheel

Author: oraluben

src/auditwheel/pool.py

@@ -1,8 +1,14 @@
+import functools
 from concurrent.futures import Future, ThreadPoolExecutor
 from pathlib import Path
+from time import sleep
 from typing import Any, Callable, Optional
 
 
+def yield_thread() -> None:
+    sleep(0)
+
+
 class FileTaskExecutor:
     """A task executor that manages concurrent file operations with deduplication.
 
@@ -28,20 +34,56 @@ def __init__(self, concurrent: int = 0):
         )
         self.working_map: dict[Path, Future[Any]] = {}
 
+    def submit_chain(
+        self, path: Path, fn: Callable[..., Any], /, *args: Any, **kwargs: Any
+    ) -> Future[Any]:
+        return self._submit(path, fn, True, *args, **kwargs)
+
     def submit(
         self, path: Path, fn: Callable[..., Any], /, *args: Any, **kwargs: Any
-    ) -> None:
+    ) -> Future[Any]:
+        return self._submit(path, fn, False, *args, **kwargs)
+
+    def _submit(
+        self,
+        path: Path,
+        fn: Callable[..., Any],
+        chain: bool,
+        /,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Future[Any]:
         if not path.is_absolute():
             path = path.absolute()
 
         future: Future[Any]
         if self.executor is None:
-            fn(*args, **kwargs)
-            return
+            future = Future()
+            future.set_result(fn(*args, **kwargs))
+        elif not chain:
+            assert path not in self.working_map, "path already in working_map"
+            future = self.executor.submit(fn, *args, **kwargs)
+            self.working_map[path] = future
+        else:
+            current = self.working_map[path]
+            future = Future()
+
+            @functools.wraps(fn)
+            def new_fn(_current: Future[Any]) -> None:
+                nonlocal future, current
+
+                assert _current == current
+
+                self.working_map.pop(path)
+                self.working_map[path] = future
+                try:
+                    future.set_result(fn(*args, **kwargs))
+                except Exception as e:
+                    future.set_exception(e)
+
+            current.add_done_callback(new_fn)
 
-        assert path not in self.working_map, "path already in working_map"
-        future = self.executor.submit(fn, *args, **kwargs)
-        self.working_map[path] = future
+        return future
 
     def wait(self, path: Optional[Path] = None) -> None:
         """Wait for tasks to complete.
@@ -55,12 +97,16 @@ def wait(self, path: Optional[Path] = None) -> None:
         """
         if self.executor is None:
             return
-        if path is None:
-            for future in self.working_map.values():
-                future.result()
-            self.working_map.clear()
-        elif path in self.working_map:
-            self.working_map.pop(path).result()
+        if path is not None:
+            while True:
+                yield_thread()
+                if path not in self.working_map:
+                    return
+                self.working_map.pop(path).result()
+        else:
+            while self.working_map:
+                path = next(iter(self.working_map))
+                self.wait(path)
 
     def __contains__(self, fn: Path) -> bool:
         return self.executor is not None and fn in self.working_map

src/auditwheel/repair.py

@@ -14,11 +14,11 @@
 from subprocess import check_call
 
 from auditwheel.patcher import ElfPatcher
-from auditwheel.pool import POOL
 
 from .elfutils import elf_read_dt_needed, elf_read_rpaths, is_subdir
 from .hashfile import hashfile
 from .policy import WheelPolicies, get_replace_platforms
+from .pool import POOL
 from .wheel_abi import get_wheel_elfdata
 from .wheeltools import InWheelCtx, add_platforms
 
@@ -83,25 +83,32 @@ def repair_wheel(
 
                 if not dest_dir.exists():
                     dest_dir.mkdir()
-                new_soname, new_path = copylib(src_path, dest_dir, patcher, dry=True)
-                if new_path not in POOL:
-                    POOL.submit(new_path, copylib, src_path, dest_dir, patcher)
-                soname_map[soname] = (new_soname, new_path)
-                replacements.append((soname, new_soname))
+                new_soname, new_path = get_new_soname(src_path, dest_dir)
+                if soname not in soname_map:
+                    soname_map[soname] = (new_soname, new_path)
+                    replacements.append((soname, new_soname))
 
-            POOL.wait()
+                    POOL.submit(new_path, copylib, src_path, dest_dir, patcher)
 
             if replacements:
-                patcher.replace_needed(fn, *replacements)
+                POOL.submit(fn, patcher.replace_needed, fn, *replacements)
 
             if len(ext_libs) > 0:
-                new_fn = fn
-                if _path_is_script(fn):
-                    new_fn = _replace_elf_script_with_shim(match.group("name"), fn)
 
-                new_rpath = os.path.relpath(dest_dir, new_fn.parent)
-                new_rpath = os.path.join("$ORIGIN", new_rpath)
-                append_rpath_within_wheel(new_fn, new_rpath, ctx.name, patcher)
+                def _patch_fn(fn: Path) -> None:
+                    new_fn = fn
+                    if _path_is_script(fn):
+                        POOL.wait(fn)
+                        new_fn = _replace_elf_script_with_shim(match.group("name"), fn)
+
+                    new_rpath = os.path.relpath(dest_dir, new_fn.parent)
+                    new_rpath = os.path.join("$ORIGIN", new_rpath)
+
+                    append_rpath_within_wheel(new_fn, new_rpath, ctx.name, patcher)
+
+                POOL.submit_chain(fn, _patch_fn)
+
+        POOL.wait()
 
         # we grafted in a bunch of libraries and modified their sonames, but
         # they may have internal dependencies (DT_NEEDED) on one another, so
@@ -133,9 +140,21 @@ def strip_symbols(libraries: Iterable[Path]) -> None:
         check_call(["strip", "-s", lib])
 
 
-def copylib(
-    src_path: Path, dest_dir: Path, patcher: ElfPatcher, dry: bool = False
-) -> tuple[str, Path]:
+def get_new_soname(src_path: Path, dest_dir: Path) -> tuple[str, Path]:
+    with open(src_path, "rb") as f:
+        shorthash = hashfile(f)[:8]
+    src_name = src_path.name
+    base, ext = src_name.split(".", 1)
+    if not base.endswith(f"-{shorthash}"):
+        new_soname = f"{base}-{shorthash}.{ext}"
+    else:
+        new_soname = src_name
+
+    dest_path = dest_dir / new_soname
+    return new_soname, dest_path
+
+
+def copylib(src_path: Path, dest_dir: Path, patcher: ElfPatcher) -> None:
     """Graft a shared library from the system into the wheel and update the
     relevant links.
 
@@ -148,19 +167,10 @@ def copylib(
     # if the library has a RUNPATH/RPATH we clear it and set RPATH to point to
     # its new location.
 
-    with open(src_path, "rb") as f:
-        shorthash = hashfile(f)[:8]
+    new_soname, dest_path = get_new_soname(src_path, dest_dir)
 
-    src_name = src_path.name
-    base, ext = src_name.split(".", 1)
-    if not base.endswith(f"-{shorthash}"):
-        new_soname = f"{base}-{shorthash}.{ext}"
-    else:
-        new_soname = src_name
-
-    dest_path = dest_dir / new_soname
-    if dry or dest_path.exists():
-        return new_soname, dest_path
+    if dest_path.exists():
+        return
 
     logger.debug("Grafting: %s -> %s", src_path, dest_path)
     rpaths = elf_read_rpaths(src_path)
@@ -174,8 +184,6 @@ def copylib(
     if any(itertools.chain(rpaths["rpaths"], rpaths["runpaths"])):
         patcher.set_rpath(dest_path, "$ORIGIN")
 
-    return new_soname, dest_path
-
 
 def append_rpath_within_wheel(
     lib_name: Path, rpath: str, wheel_base_dir: Path, patcher: ElfPatcher