chore(lib-injection): unpack wheels to arch specific directories

lib-injection/dl_wheels.py

@@ -127,33 +127,36 @@
         if not args.dry_run:
             subprocess.run(cmd, capture_output=not args.verbose, check=True)
 
-    wheel_files = [f for f in os.listdir(dl_dir) if f.endswith(".whl")]
-    for whl in wheel_files:
-        wheel_file = os.path.join(dl_dir, whl)
-        print("Unpacking %s" % wheel_file)
-        # -q for quieter output, else we get all the files being unzipped.
-        subprocess.run(
-            [
-                "unzip",
-                "-q",
-                "-o",
-                wheel_file,
-                "-d",
-                os.path.join(dl_dir, "site-packages-ddtrace-py%s-%s" % (python_version, platform)),
-            ]
-        )
-        # Remove the wheel as it has been unpacked
-        os.remove(wheel_file)
-
-    sitepackages_root = Path(dl_dir) / f"site-packages-ddtrace-py{python_version}-{platform}"
-    directories_to_remove = [
-        sitepackages_root / "google" / "protobuf",
-        sitepackages_root / "google" / "_upb",
-    ]
-    directories_to_remove.extend(sitepackages_root.glob("protobuf-*"))  # dist-info directories
-
-    for directory in directories_to_remove:
-        try:
-            shutil.rmtree(directory)
-        except Exception:
-            pass
+        # FIXED: Move wheel unpacking inside the architecture loop to prevent overwrites
+        # Create architecture-specific directory name
+        arch_platform_dir = "site-packages-ddtrace-py%s-%s-%s" % (python_version, platform, arch)
+        wheel_files = [f for f in os.listdir(dl_dir) if f.endswith(".whl")]
+        for whl in wheel_files:
+            wheel_file = os.path.join(dl_dir, whl)
+            print("Unpacking %s to %s" % (wheel_file, arch_platform_dir))
+            # -q for quieter output, else we get all the files being unzipped.
+            subprocess.run(
+                [
+                    "unzip",
+                    "-q",
+                    "-o",
+                    wheel_file,
+                    "-d",
+                    os.path.join(dl_dir, arch_platform_dir),
+                ]
+            )
+            # Remove the wheel as it has been unpacked
+            os.remove(wheel_file)
+
+        sitepackages_root = Path(dl_dir) / arch_platform_dir
+        directories_to_remove = [
+            sitepackages_root / "google" / "protobuf",
+            sitepackages_root / "google" / "_upb",
+        ]
+        directories_to_remove.extend(sitepackages_root.glob("protobuf-*"))  # dist-info directories
+
+        for directory in directories_to_remove:
+            try:
+                shutil.rmtree(directory)
+            except Exception:
+                pass

lib-injection/sources/sitecustomize.py

@@ -348,7 +348,15 @@ def _inject():
         _log("user-installed ddtrace not found, configuring application to use injection site-packages")
 
         current_platform = "manylinux2014" if _get_clib() == "gnu" else "musllinux_1_2"
-        _log("detected platform %s" % current_platform, level="debug")
+        # Determine architecture
+        current_arch = platform.machine()
+        if current_arch == "x86_64":
+            current_arch = "x86_64"
+        elif current_arch in ["aarch64", "arm64"]:
+            current_arch = "aarch64"
+        else:
+            current_arch = "x86_64"  # Default fallback
+        _log("detected platform %s, architecture %s" % (current_platform, current_arch), level="debug")
 
         pkgs_path = os.path.join(SCRIPT_DIR, "ddtrace_pkgs")
         _log("ddtrace_pkgs path is %r" % pkgs_path, level="debug")
@@ -462,9 +470,23 @@ def _inject():
             )
             return
 
-        site_pkgs_path = os.path.join(
+        # Try architecture-specific directory first (new format), fall back to platform-only (legacy)
+        arch_site_pkgs_path = os.path.join(
+            pkgs_path,
+            "%s%s-%s-%s" % (SITE_PKGS_MARKER, ".".join(PYTHON_VERSION.split(".")[:2]), current_platform, current_arch),
+        )
+        legacy_site_pkgs_path = os.path.join(
             pkgs_path, "%s%s-%s" % (SITE_PKGS_MARKER, ".".join(PYTHON_VERSION.split(".")[:2]), current_platform)
         )
+
+        if os.path.exists(arch_site_pkgs_path):
+            site_pkgs_path = arch_site_pkgs_path
+            _log("using architecture-specific site-packages path: %r" % site_pkgs_path, level="debug")
+        elif os.path.exists(legacy_site_pkgs_path):
+            site_pkgs_path = legacy_site_pkgs_path
+            _log("using legacy site-packages path: %r" % site_pkgs_path, level="debug")
+        else:
+            site_pkgs_path = arch_site_pkgs_path  # Use new format for error reporting
         _log("site-packages path is %r" % site_pkgs_path, level="debug")
         if not os.path.exists(site_pkgs_path):
             _log("ddtrace site-packages not found in %r, aborting" % site_pkgs_path, level="error")

tests/lib_injection/conftest.py

@@ -37,7 +37,16 @@ def get_platform_details():
     python_version = platform.python_version()
     py_major_minor = ".".join(python_version.split(".")[:2])
 
-    return py_major_minor, "manylinux2014"
+    # Determine architecture
+    current_arch = platform.machine()
+    if current_arch == "x86_64":
+        current_arch = "x86_64"
+    elif current_arch in ["aarch64", "arm64"]:
+        current_arch = "aarch64"
+    else:
+        current_arch = "x86_64"  # Default fallback
+
+    return py_major_minor, "manylinux2014", current_arch
 
 
 @pytest.fixture(scope="session")
@@ -61,8 +70,8 @@ def ddtrace_injection_artifact():
         )
 
         # 2. Create the target site-packages directory structure
-        py_major_minor, platform_tag = get_platform_details()
-        target_site_packages_name = f"site-packages-ddtrace-py{py_major_minor}-{platform_tag}"
+        py_major_minor, platform_tag, arch = get_platform_details()
+        target_site_packages_name = f"site-packages-ddtrace-py{py_major_minor}-{platform_tag}-{arch}"
         target_site_packages_path = os.path.join(sources_dir_in_session_tmp, "ddtrace_pkgs", target_site_packages_name)
         os.makedirs(target_site_packages_path, exist_ok=True)