[compiler-rt] [XRay][test] Pin CPU affinity in fdr-mode.cpp to deflake UNWRITE checks (PR #211595)

[Nikita Taranov via llvm-commits]

https://github.com/nickitat created https://github.com/llvm/llvm-project/pull/211595

The UNWRITE portion of fdr-mode.cpp asserts that, with a high func_duration_threshold_us, all short function records are "unwritten" (erased) except arg1-logging records. FDR mode only rewinds a function's enter record when its exit lands within the threshold AND no metadata record was written in between: XRayFDRController::functionExit rewinds only when recordPreamble() returns NoChange. A CPU migration between the enter and exit makes recordPreamble() emit a NewCPUId metadata record and return WroteMetadata, which blocks the rewind and leaves stray function-enter/exit records, failing the UNWRITE-NEXT/UNWRITE-NOT checks.

On busy multi-core machines this migration is likely enough to flake the test ([observed](https://github.com/llvm/llvm-project/pull/210739#issuecomment-5053841943) on the sanitizer-aarch64-linux buildbot). The test already documented the hazard ("A CPU switch could erroneously fail this test").

Pin the process to a single CPU at the start of main() so no migration can occur; child threads inherit the affinity mask. Guarded to Linux; a no-op on other platforms.

>From 96e6b8804b3c132c0302701977a459b950ed3133 Mon Sep 17 00:00:00 2001
From: Nikita Taranov <nikita.taranov at clickhouse.com>
Date: Thu, 23 Jul 2026 22:05:39 +0700
Subject: [PATCH] [XRay][test] Pin CPU affinity in fdr-mode.cpp to deflake
 UNWRITE checks

The UNWRITE portion of fdr-mode.cpp asserts that, with a high
func_duration_threshold_us, all short function records are "unwritten"
(erased) except arg1-logging records. FDR mode only rewinds a function's
enter record when its exit lands within the threshold AND no metadata
record was written in between: XRayFDRController::functionExit rewinds
only when recordPreamble() returns NoChange. A CPU migration between the
enter and exit makes recordPreamble() emit a NewCPUId metadata record and
return WroteMetadata, which blocks the rewind and leaves stray
function-enter/exit records, failing the UNWRITE-NEXT/UNWRITE-NOT checks.

On busy multi-core machines this migration is likely enough to flake the
test (observed on the sanitizer-aarch64-linux buildbot). The test already
documented the hazard ("A CPU switch could erroneously fail this test").

Pin the process to a single CPU at the start of main() so no migration can
occur; child threads inherit the affinity mask. The CPU is chosen from the
process's allowed affinity set rather than assuming CPU 0 is usable, and
the pinning degrades gracefully (no-op) if the affinity calls fail. Guarded
to Linux; a no-op on other platforms.
---
 .../test/xray/TestCases/Posix/fdr-mode.cpp    | 39 ++++++++++++++++++-
 1 file changed, 37 insertions(+), 2 deletions(-)

diff --git a/compiler-rt/test/xray/TestCases/Posix/fdr-mode.cpp b/compiler-rt/test/xray/TestCases/Posix/fdr-mode.cpp
index 08152d3179b17..071e5b0319521 100644
--- a/compiler-rt/test/xray/TestCases/Posix/fdr-mode.cpp
+++ b/compiler-rt/test/xray/TestCases/Posix/fdr-mode.cpp
@@ -28,6 +28,15 @@
 // UNSUPPORTED: armhf-linux
 // REQUIRES: built-in-llvm-tree
 
+#if defined(__linux__)
+// Pull in sched_setaffinity/cpu_set_t (GNU extensions). _GNU_SOURCE must be
+// defined before any libc header is included.
+#ifndef _GNU_SOURCE
+#define _GNU_SOURCE
+#endif
+#include <sched.h>
+#endif
+
 #include "xray/xray_log_interface.h"
 #include <cassert>
 #include <chrono>
@@ -48,6 +57,29 @@ thread_local uint64_t var = 0;
 void __attribute__((noinline)) fArg(int) { }
 
 int main(int argc, char *argv[]) {
+#if defined(__linux__)
+  // Pin to a single CPU. The UNWRITE portion of this test relies on FDR mode
+  // "unwriting" (rewinding) short function records, which only happens when a
+  // function's exit record immediately follows its enter record in the
+  // per-thread buffer. A CPU migration in between forces a NewCPUId metadata
+  // record, which blocks the rewind and leaves stray function-enter/exit
+  // records, flaking the UNWRITE checks below. Pinning prevents the migration;
+  // child threads inherit this affinity mask. Pick the first CPU from the
+  // process's allowed set rather than assuming CPU 0 is usable.
+  cpu_set_t Allowed;
+  CPU_ZERO(&Allowed);
+  if (sched_getaffinity(0, sizeof(Allowed), &Allowed) == 0) {
+    for (int Cpu = 0; Cpu < CPU_SETSIZE; ++Cpu) {
+      if (CPU_ISSET(Cpu, &Allowed)) {
+        cpu_set_t Pinned;
+        CPU_ZERO(&Pinned);
+        CPU_SET(Cpu, &Pinned);
+        (void)sched_setaffinity(0, sizeof(Pinned), &Pinned);
+        break;
+      }
+    }
+  }
+#endif
   std::cout << "Logging before init." << std::endl;
   // CHECK: Logging before init.
   assert(__xray_log_select_mode("xray-fdr") ==
@@ -106,8 +138,11 @@ int main(int argc, char *argv[]) {
 // TRACE-DAG: - { type: 0, func-id: [[FIDARG:[0-9]+]], function: 'fArg(int)', args: [ 1 ], cpu: {{.*}}, thread: [[THREAD2]], process: [[PROCESS]], kind: function-enter-arg, tsc: {{[0-9]+}}, data: '' }
 // TRACE-DAG: - { type: 0, func-id: [[FIDARG]], function: 'fArg(int)', cpu: {{.*}}, thread: [[THREAD2]], process: [[PROCESS]], kind: function-exit, tsc: {{[0-9]+}}, data: '' }
 
-// Assert that when unwriting is enabled with a high threshold time, all the function records are erased. A CPU switch could erroneously fail this test, but
-// is unlikely given the test program.
+// Assert that when unwriting is enabled with a high threshold time, all the
+// function records are erased. A CPU switch between a function's enter and exit
+// would block the unwrite and fail this test; on Linux main() pins the process
+// to a single CPU to prevent that (elsewhere it remains theoretically possible
+// but unlikely given the test program).
 // Even with a high threshold, arg1 logging is never unwritten.
 // UNWRITE: header:
 // UNWRITE: records: