parallel get_wheel_elfdata

Author: oraluben

src/auditwheel/pool.py

@@ -1,4 +1,3 @@
-import time
 from concurrent.futures import Future, ThreadPoolExecutor
 from pathlib import Path
 from typing import Any, Callable, Optional
@@ -21,16 +20,16 @@ class FileTaskExecutor:
         >>> executor.wait()  # Wait for all tasks to complete
     """
 
-    def __init__(self, concurrent: int = 1):
+    def __init__(self, concurrent: int = 0):
         self.executor = (
             None
             if concurrent == 1
             else ThreadPoolExecutor(concurrent if concurrent > 1 else None)
         )
-        self.working_map: dict[Path, Future[tuple[str, str]]] = {}
+        self.working_map: dict[Path, Future[Any]] = {}
 
     def submit(
-        self, path: Path, fn: Callable[[Any], Any], /, *args: Any, **kwargs: Any
+        self, path: Path, fn: Callable[..., Any], /, *args: Any, **kwargs: Any
     ) -> None:
         if not path.is_absolute():
             path = path.absolute()
@@ -39,12 +38,11 @@ def submit(
         if self.executor is None:
             future = Future()
             future.set_result(fn(*args, **kwargs))
-            return None
+            return
 
         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.
@@ -62,18 +60,8 @@ def wait(self, path: Optional[Path] = None) -> None:
             for future in self.working_map.values():
                 future.result()
             self.working_map.clear()
-        elif future := self.working_map.pop(path, None):
-            future.result()
-
-
-def fake_job(i: int) -> int:
-    print(f"start {i}")
-    time.sleep(i)
-    print(f"end {i}")
+        elif path in self.working_map:
+            self.working_map.pop(path).result()
 
 
-if __name__ == "__main__":
-    executor = FileTaskExecutor(concurrent=0)
-    for i in range(10):
-        executor.submit(Path(f"test{i}.txt"), fake_job, i)
-    executor.wait()
+POOL = FileTaskExecutor()

src/auditwheel/wheel_abi.py

@@ -10,6 +10,10 @@
 from pathlib import Path
 from typing import Optional, TypeVar
 
+from elftools.elf.elffile import ELFFile
+
+from auditwheel.pool import POOL
+
 from . import json
 from .architecture import Architecture
 from .elfutils import (
@@ -94,6 +98,55 @@ def get_wheel_elfdata(
         shared_libraries_with_invalid_machine = []
 
         platform_wheel = False
+
+        def inner(fn: Path, elf: ELFFile) -> None:
+            nonlocal \
+                platform_wheel, \
+                shared_libraries_in_purelib, \
+                uses_ucs2_symbols, \
+                uses_PyFPE_jbuf
+
+            so_name = fn.name
+
+            log.debug("processing: %s", fn)
+            elftree = ldd(fn, exclude=exclude)
+
+            try:
+                arch = elftree.platform.baseline_architecture
+                if arch != wheel_policy.architecture.baseline:
+                    shared_libraries_with_invalid_machine.append(so_name)
+                    log.warning("ignoring: %s with %s architecture", so_name, arch)
+                    return
+            except ValueError:
+                shared_libraries_with_invalid_machine.append(so_name)
+                log.warning("ignoring: %s with unknown architecture", so_name)
+                return
+
+            platform_wheel = True
+
+            for key, value in elf_find_versioned_symbols(elf):
+                log.debug("key %s, value %s", key, value)
+                versioned_symbols[key].add(value)
+
+            is_py_ext, py_ver = elf_is_python_extension(fn, elf)
+
+            # If the ELF is a Python extention, we definitely need to
+            # include its external dependencies.
+            if is_py_ext:
+                full_elftree[fn] = elftree
+                uses_PyFPE_jbuf |= elf_references_PyFPE_jbuf(elf)
+                if py_ver == 2:
+                    uses_ucs2_symbols |= any(True for _ in elf_find_ucs2_symbols(elf))
+                full_external_refs[fn] = wheel_policy.lddtree_external_references(
+                    elftree, ctx.path
+                )
+            else:
+                # If the ELF is not a Python extension, it might be
+                # included in the wheel already because auditwheel repair
+                # vendored it, so we will check whether we should include
+                # its internal references later.
+                nonpy_elftree[fn] = elftree
+
         for fn, elf in elf_file_filter(ctx.iter_files()):
             # Check for invalid binary wheel format: no shared library should
             # be found in purelib
@@ -104,49 +157,8 @@ def get_wheel_elfdata(
             if any(p.name == "purelib" for p in fn.parents):
                 shared_libraries_in_purelib.append(so_name)
 
-            # If at least one shared library exists in purelib, this is going
-            # to fail and there's no need to do further checks
             if not shared_libraries_in_purelib:
-                log.debug("processing: %s", fn)
-                elftree = ldd(fn, exclude=exclude)
-
-                try:
-                    arch = elftree.platform.baseline_architecture
-                    if arch != wheel_policy.architecture.baseline:
-                        shared_libraries_with_invalid_machine.append(so_name)
-                        log.warning("ignoring: %s with %s architecture", so_name, arch)
-                        continue
-                except ValueError:
-                    shared_libraries_with_invalid_machine.append(so_name)
-                    log.warning("ignoring: %s with unknown architecture", so_name)
-                    continue
-
-                platform_wheel = True
-
-                for key, value in elf_find_versioned_symbols(elf):
-                    log.debug("key %s, value %s", key, value)
-                    versioned_symbols[key].add(value)
-
-                is_py_ext, py_ver = elf_is_python_extension(fn, elf)
-
-                # If the ELF is a Python extention, we definitely need to
-                # include its external dependencies.
-                if is_py_ext:
-                    full_elftree[fn] = elftree
-                    uses_PyFPE_jbuf |= elf_references_PyFPE_jbuf(elf)
-                    if py_ver == 2:
-                        uses_ucs2_symbols |= any(
-                            True for _ in elf_find_ucs2_symbols(elf)
-                        )
-                    full_external_refs[fn] = wheel_policy.lddtree_external_references(
-                        elftree, ctx.path
-                    )
-                else:
-                    # If the ELF is not a Python extension, it might be
-                    # included in the wheel already because auditwheel repair
-                    # vendored it, so we will check whether we should include
-                    # its internal references later.
-                    nonpy_elftree[fn] = elftree
+                POOL.submit(fn, inner, fn, elf)
 
         # If at least one shared library exists in purelib, raise an error
         if shared_libraries_in_purelib:
@@ -164,6 +176,8 @@ def get_wheel_elfdata(
                 wheel_policy.architecture, shared_libraries_with_invalid_machine
             )
 
+        POOL.wait()
+
         # Get a list of all external libraries needed by ELFs in the wheel.
         needed_libs = {
             lib