gh-141592: test_perf_profiler: fix race in perf map cleanup using addCleanup

Lib/test/test_perf_profiler.py

@@ -11,7 +11,7 @@
     assert_python_failure,
     assert_python_ok,
 )
-from test.support.os_helper import temp_dir
+from test.support.os_helper import temp_dir, unlink
 
 
 if not support.has_subprocess_support:
@@ -35,17 +35,20 @@ def supports_trampoline_profiling():
 
 
 class TestPerfTrampoline(unittest.TestCase):
+
     def setUp(self):
-        super().setUp()
-        self.perf_files = set(pathlib.Path("/tmp/").glob("perf-*.map"))
+        # Register cleanup for JIT files for all tests in this class
+        self.addCleanup(self._cleanup_jit_files)
 
-    def tearDown(self) -> None:
-        super().tearDown()
-        files_to_delete = (
-            set(pathlib.Path("/tmp/").glob("perf-*.map")) - self.perf_files
-        )
-        for file in files_to_delete:
-            file.unlink()
+    def _cleanup_perf_map(self, pid):
+        """Helper to safely delete a specific perf map file."""
+        unlink(f"/tmp/perf-{pid}.map")
+
+    def _cleanup_jit_files(self):
+        """Helper to safely delete JIT-related temp files."""
+        for pattern in ("jit*.dump", "jitted-*.so"):
+            for path in pathlib.Path("/tmp/").glob(pattern):
+                unlink(path)
 
     @unittest.skipIf(support.check_bolt_optimized(), "fails on BOLT instrumented binaries")
     def test_trampoline_works(self):
@@ -71,6 +74,7 @@ def baz():
                 stdout=subprocess.PIPE,
                 env=env,
             ) as process:
+                self.addCleanup(self._cleanup_perf_map, process.pid)
                 stdout, stderr = process.communicate()
 
         self.assertEqual(stderr, "")
@@ -140,11 +144,15 @@ def baz():
                 stdout=subprocess.PIPE,
                 env=env,
             ) as process:
+                self.addCleanup(self._cleanup_perf_map, process.pid)
                 stdout, stderr = process.communicate()
 
         self.assertEqual(process.returncode, 0)
         self.assertEqual(stderr, "")
+
         child_pid = int(stdout.strip())
+        self.addCleanup(self._cleanup_perf_map, child_pid)
+
         perf_file = pathlib.Path(f"/tmp/perf-{process.pid}.map")
         perf_child_file = pathlib.Path(f"/tmp/perf-{child_pid}.map")
         self.assertTrue(perf_file.exists())
@@ -199,6 +207,7 @@ def baz():
                     stdout=subprocess.PIPE,
                     env=env,
                 ) as process:
+                    self.addCleanup(self._cleanup_perf_map, process.pid)
                     stdout, stderr = process.communicate()
 
             self.assertEqual(stderr, "")
@@ -358,7 +367,6 @@ def run_perf(cwd, *args, use_jit=False, **env_vars):
         if proc.returncode:
             print(proc.stderr, file=sys.stderr)
             raise ValueError(f"Perf failed with return code {proc.returncode}")
-        # Copy the jit_output_file to the output_file
         os.rename(jit_output_file, output_file)
 
     base_cmd = ("perf", "script")
@@ -374,6 +382,10 @@ def run_perf(cwd, *args, use_jit=False, **env_vars):
 
 
 class TestPerfProfilerMixin:
+    def setUp(self):
+        # This setUp is added to make super() calls in child classes work smoothly.
+        pass
+
     def run_perf(self, script_dir, perf_mode, script):
         raise NotImplementedError()
 
@@ -433,23 +445,16 @@ def baz(n):
     is_unwinding_reliable_with_frame_pointers(),
     "Unwinding is unreliable with frame pointers",
 )
-class TestPerfProfiler(unittest.TestCase, TestPerfProfilerMixin):
-    def run_perf(self, script_dir, script, activate_trampoline=True):
-        if activate_trampoline:
-            return run_perf(script_dir, sys.executable, "-Xperf", script)
-        return run_perf(script_dir, sys.executable, script)
+class TestPerfProfiler(TestPerfTrampoline, TestPerfProfilerMixin):
 
     def setUp(self):
         super().setUp()
-        self.perf_files = set(pathlib.Path("/tmp/").glob("perf-*.map"))
+        TestPerfProfilerMixin.setUp(self)
 
-    def tearDown(self) -> None:
-        super().tearDown()
-        files_to_delete = (
-            set(pathlib.Path("/tmp/").glob("perf-*.map")) - self.perf_files
-        )
-        for file in files_to_delete:
-            file.unlink()
+    def run_perf(self, script_dir, script, activate_trampoline=True):
+        if activate_trampoline:
+            return run_perf(script_dir, sys.executable, "-Xperf", script)
+        return run_perf(script_dir, sys.executable, script)
 
     def test_pre_fork_compile(self):
         code = """if 1:
@@ -499,11 +504,14 @@ def compile_trampolines_for_all_functions():
                 stdout=subprocess.PIPE,
                 env=env,
             ) as process:
+                self.addCleanup(self._cleanup_perf_map, process.pid)
                 stdout, stderr = process.communicate()
 
         self.assertEqual(process.returncode, 0)
         self.assertNotIn("Error:", stderr)
         child_pid = int(stdout.strip())
+        self.addCleanup(self._cleanup_perf_map, child_pid)
+
         perf_file = pathlib.Path(f"/tmp/perf-{process.pid}.map")
         perf_child_file = pathlib.Path(f"/tmp/perf-{child_pid}.map")
         self.assertTrue(perf_file.exists())
@@ -549,27 +557,21 @@ def _is_perf_version_at_least(major, minor):
 @unittest.skipUnless(
     _is_perf_version_at_least(6, 6), "perf command may not work due to a perf bug"
 )
-class TestPerfProfilerWithDwarf(unittest.TestCase, TestPerfProfilerMixin):
+class TestPerfProfilerWithDwarf(TestPerfTrampoline, TestPerfProfilerMixin):
+
+    def setUp(self):
+        # This calls TestPerfTrampoline.setUp() which registers jit_cleanup
+        super().setUp()
+        # This calls TestPerfProfilerMixin.setUp()
+        TestPerfProfilerMixin.setUp(self)
+
     def run_perf(self, script_dir, script, activate_trampoline=True):
         if activate_trampoline:
             return run_perf(
                 script_dir, sys.executable, "-Xperf_jit", script, use_jit=True
             )
         return run_perf(script_dir, sys.executable, script, use_jit=True)
 
-    def setUp(self):
-        super().setUp()
-        self.perf_files = set(pathlib.Path("/tmp/").glob("jit*.dump"))
-        self.perf_files |= set(pathlib.Path("/tmp/").glob("jitted-*.so"))
-
-    def tearDown(self) -> None:
-        super().tearDown()
-        files_to_delete = set(pathlib.Path("/tmp/").glob("jit*.dump"))
-        files_to_delete |= set(pathlib.Path("/tmp/").glob("jitted-*.so"))
-        files_to_delete = files_to_delete - self.perf_files
-        for file in files_to_delete:
-            file.unlink()
-
 
 if __name__ == "__main__":
     unittest.main()