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,30 +20,28 @@ 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()
 
         future: Future[Any]
         if self.executor is None:
-            future = Future()
-            future.set_result(fn(*args, **kwargs))
-            return None
+            fn(*args, **kwargs)
+            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 +59,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(2)

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,19 +98,19 @@ def get_wheel_elfdata(
         shared_libraries_with_invalid_machine = []
 
         platform_wheel = False
-        for fn, elf in elf_file_filter(ctx.iter_files()):
-            # Check for invalid binary wheel format: no shared library should
-            # be found in purelib
-            so_name = fn.name
 
-            # If this is in purelib, add it to the list of shared libraries in
-            # purelib
-            if any(p.name == "purelib" for p in fn.parents):
-                shared_libraries_in_purelib.append(so_name)
+        def inner(fn: Path) -> None:
+            nonlocal \
+                platform_wheel, \
+                shared_libraries_in_purelib, \
+                uses_ucs2_symbols, \
+                uses_PyFPE_jbuf
+
+            with open(fn, "rb") as f:
+                elf = ELFFile(f)
+
+                so_name = fn.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)
 
@@ -115,11 +119,11 @@ def get_wheel_elfdata(
                     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
+                        return
                 except ValueError:
                     shared_libraries_with_invalid_machine.append(so_name)
                     log.warning("ignoring: %s with unknown architecture", so_name)
-                    continue
+                    return
 
                 platform_wheel = True
 
@@ -148,6 +152,20 @@ def get_wheel_elfdata(
                     # its internal references later.
                     nonpy_elftree[fn] = elftree
 
+        # Create new ELFFile object to avoid use-after-free
+        for fn, _elf in elf_file_filter(ctx.iter_files()):
+            # Check for invalid binary wheel format: no shared library should
+            # be found in purelib
+            so_name = fn.name
+
+            # If this is in purelib, add it to the list of shared libraries in
+            # purelib
+            if any(p.name == "purelib" for p in fn.parents):
+                shared_libraries_in_purelib.append(so_name)
+
+            if not shared_libraries_in_purelib:
+                POOL.submit(fn, inner, fn)
+
         # If at least one shared library exists in purelib, raise an error
         if shared_libraries_in_purelib:
             libraries = "\n\t".join(shared_libraries_in_purelib)
@@ -159,6 +177,8 @@ def get_wheel_elfdata(
             )
             raise RuntimeError(msg)
 
+        POOL.wait()
+
         if not platform_wheel:
             raise NonPlatformWheel(
                 wheel_policy.architecture, shared_libraries_with_invalid_machine