test(profiling): unflake some tests

tests/profiling/collector/pprof_utils.py

@@ -143,7 +143,9 @@ def __init__(self, *args, **kwargs):
         super().__init__(event_type=LockEventType.RELEASE, *args, **kwargs)
 
 
-def parse_newest_profile(filename_prefix: str, assert_samples: bool = True) -> pprof_pb2.Profile:
+def parse_newest_profile(
+    filename_prefix: str, assert_samples: bool = True, allow_penultimate: bool = False
+) -> pprof_pb2.Profile:
     """Parse the newest profile that has given filename prefix. The profiler
     outputs profile file with following naming convention:
     <filename_prefix>.<pid>.<counter>.pprof, and in tests, we'd want to parse
@@ -160,6 +162,18 @@ def parse_newest_profile(filename_prefix: str, assert_samples: bool = True) -> p
     profile = pprof_pb2.Profile()
     profile.ParseFromString(serialized_data)
 
+    if allow_penultimate and len(profile.sample) == 0 and len(files) > 1:
+        # The newest profile file may be empty if it was just created and has not yet accumulated samples.
+        # In this case, we try to parse the (second-to-last) file instead, which is more likely
+        # to contain complete data. We temporarily rename the newest file so parse_newest_profile picks
+        # the previous one.
+        temp_filename = filename + ".temp"
+        os.rename(filename, temp_filename)
+        try:
+            return parse_newest_profile(filename_prefix, assert_samples, allow_penultimate=False)
+        finally:
+            os.rename(temp_filename, filename)
+
     if assert_samples:
         assert len(profile.sample) > 0, "No samples found in profile"
 

tests/profiling/test_main.py

@@ -60,7 +60,7 @@ def test_call_script_pprof_output(tmp_path: Path) -> None:
 
     stdout = stdout.decode() if isinstance(stdout, bytes) else stdout
     _, pid = list(s.strip() for s in stdout.strip().split("\n"))
-    profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}")
+    profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}", allow_penultimate=True)
     samples = pprof_utils.get_samples_with_value_type(profile, "cpu-time")
     assert len(samples) > 0
 
@@ -77,7 +77,7 @@ def test_fork(tmp_path: Path) -> None:
     assert exitcode == 0
     stdout = stdout.decode() if isinstance(stdout, bytes) else stdout
     child_pid = stdout.strip()
-    profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}")
+    profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}", allow_penultimate=True)
     parent_expected_acquire_events = [
         pprof_utils.LockAcquireEvent(
             caller_name="<module>",
@@ -99,7 +99,7 @@ def test_fork(tmp_path: Path) -> None:
         expected_acquire_events=parent_expected_acquire_events,
         expected_release_events=parent_expected_release_events,
     )
-    child_profile = pprof_utils.parse_newest_profile(filename + "." + str(child_pid))
+    child_profile = pprof_utils.parse_newest_profile(f"{filename}.{child_pid}")
     # We expect the child profile to not have lock events from the parent process
     # Note that assert_lock_events function only checks that the given events
     # exists, and doesn't assert that other events don't exist.