[Support] Optimize signal handling file removal code (#173586)

clangd/clangd#1787 describes how LLVM's code for removing partially
written files when taking a fatal signal is slower than it should be.
This code was substantially rewritten by JF Bastien in
aa1333a91f to make it (more) signal-safe.

As written in 2018, the logic always allocates a new node for every file
we attempt to protect, and those nodes are added to a global singly
linked list and never removed. If you open a lot of output files,
suddenly output file opening becomes O(n^2), which is what happened
during clangd indexing.

Removing files on signals in a multi-threaded environment is really
complicated! We can use locks to synchronize between threads that are
writing to the list, but we cannot use locks to synchronize against
re-entrant signals, and we don't have any great tools for masking or
delaying things like SIGTERM. This makes it difficult to assert that we
have the one and only reference to a node, even after removing it from
the linked structure, so we can safely deallocate it.

This implementation sidesteps that problem by reusing nodes on the list
with a null Filename pointer. The Filename is null if the file no longer
needs to be removed. We replace the filename with a different sentinel
while removing files using a CAS operation to simulate a lock acquisition.
* Case 1: This is the normal case, works as intended, the list should
grow in length to the number of concurrently open files that must be
removed on exit, which should be O(# threads)
* Case 2: In this case, the signal handler will leak the filename string
without a GC root for it, which is fine, since we're exiting.

I added a benchmark, and this
change makes it 16x faster, and the profile shows that the time is no
longer mostly spent in this signal handling code.

---------

Co-authored-by: JF Bastien <github@jfbastien.com>

llvm/benchmarks/CMakeLists.txt

@@ -18,6 +18,7 @@ add_benchmark(PointerUnionBM PointerUnionBM.cpp PARTIAL_SOURCES_INTENDED)
 add_benchmark(ImmutableSetIteratorBM ImmutableSetIteratorBM.cpp PARTIAL_SOURCES_INTENDED)
 
 add_benchmark(RuntimeLibcallsBench RuntimeLibcalls.cpp PARTIAL_SOURCES_INTENDED)
+add_benchmark(writeToOutputInParallelBench writeToOutputInParallel.cpp PARTIAL_SOURCES_INTENDED)
 
 if(NOT LLVM_TOOL_LLVM_DRIVER_BUILD)
   # TODO: Check if the tools are in LLVM_DISTRIBUTION_COMPONENTS with

llvm/benchmarks/writeToOutputInParallel.cpp

@@ -0,0 +1,113 @@
+//===- writeToOutputInParallel.cpp - Parallel file writing benchmark ------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "benchmark/benchmark.h"
+#include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/raw_ostream.h"
+#include <string>
+#include <thread>
+#include <vector>
+
+using namespace llvm;
+
+// Benchmark parallel file writing using writeToOutput. This simulates scenarios
+// where multiple threads are writing files concurrently, which is common in
+// parallel compilation scenarios. The goal of this benchmark is to ensure that
+// LLVM's global signal handling state updates aren't too expensive.
+static void BM_WriteToOutputInParallel(benchmark::State &State) {
+  const int NumThreads = State.range(0);
+  const int FilesPerThread = State.range(1);
+  const int BytesPerFile = 40 * 1024; // ~40KB per file
+
+  for (auto _ : State) {
+    // Create one top-level unique directory
+    SmallString<128> TopLevelDir;
+    if (sys::fs::createUniqueDirectory("writeToOutputBM", TopLevelDir)) {
+      State.SkipWithError("Failed to create temporary directory");
+      return;
+    }
+    auto Cleanup =
+        llvm::scope_exit([&]() { sys::fs::remove_directories(TopLevelDir); });
+
+    // Create subdirectories for each thread within the top-level directory
+    std::vector<SmallString<128>> ThreadDirs;
+    for (int I = 0; I < NumThreads; ++I) {
+      SmallString<128> ThreadDir(TopLevelDir);
+      sys::path::append(ThreadDir, "thread_" + std::to_string(I));
+      if (sys::fs::create_directory(ThreadDir)) {
+        State.SkipWithError("Failed to create thread directory");
+        return;
+      }
+      ThreadDirs.push_back(ThreadDir);
+    }
+
+    // Launch threads, each writing multiple files
+    std::vector<std::thread> Threads;
+    for (int ThreadIdx = 0; ThreadIdx < NumThreads; ++ThreadIdx) {
+      Threads.emplace_back([&, ThreadIdx]() {
+        const auto &ThreadDir = ThreadDirs[ThreadIdx];
+        for (int FileIdx = 0; FileIdx < FilesPerThread; ++FileIdx) {
+          SmallString<128> Path(ThreadDir);
+          sys::path::append(Path, "file_" + std::to_string(FileIdx) + ".bin");
+
+          Error E = writeToOutput(Path, [=](raw_ostream &Out) -> Error {
+            // Write 32-bit integers up to BytesPerFile
+            const int NumInts = BytesPerFile / sizeof(int32_t);
+            for (int32_t I = 0; I < NumInts; ++I) {
+              Out.write(reinterpret_cast<const char *>(&I), sizeof(I));
+            }
+            return Error::success();
+          });
+          if (E) {
+            State.SkipWithError("Failed to create outputfile " +
+                                Path.str().str());
+            return;
+          }
+        }
+      });
+    }
+
+    // Wait for all threads to complete
+    for (auto &Thread : Threads) {
+      Thread.join();
+    }
+
+    // Cleanup happens automatically via scope_exit
+  }
+
+  // Report throughput metrics
+  const int64_t TotalFiles = NumThreads * FilesPerThread;
+  const int64_t TotalBytes = TotalFiles * BytesPerFile;
+
+  State.SetItemsProcessed(State.iterations() * TotalFiles);
+  State.SetBytesProcessed(State.iterations() * TotalBytes);
+  State.counters["threads"] = NumThreads;
+  State.counters["files_per_thread"] = FilesPerThread;
+  State.counters["total_files"] = TotalFiles;
+}
+
+// Test various combinations of thread counts and files per thread
+// These represent different parallelism scenarios:
+// - Low parallelism, many files per thread (serial-like workload)
+// - High parallelism, few files per thread (highly parallel workload)
+// - Balanced scenarios
+
+BENCHMARK(BM_WriteToOutputInParallel)
+    ->Args({1, 1000})  // 1 thread, 1000 files
+    ->Args({2, 500})   // 2 threads, 500 files each
+    ->Args({4, 250})   // 4 threads, 250 files each
+    ->Args({8, 125})   // 8 threads, 125 files each
+    ->Args({10, 100})  // 10 threads, 100 files each
+    ->Args({10, 1000}) // 10 threads, 1000 files each (stress test)
+    ->Unit(benchmark::kMillisecond);
+
+BENCHMARK_MAIN();

llvm/lib/Support/Unix/Signals.inc

@@ -98,6 +98,10 @@ static std::atomic<SignalHandlerFunctionType> OneShotPipeSignalFunction =
     nullptr;
 
 namespace {
+/// Sentinel stored in a node after the signal handler has removed the file;
+/// not a valid path, never freed.
+static char InvalidPathSentinel[] = "\01\02\03\04";
+
 /// Signal-safe removal of files.
 /// Inserting and erasing from the list isn't signal-safe, but removal of files
 /// themselves is signal-safe. Memory is freed when the head is freed, deletion
@@ -107,8 +111,8 @@ class FileToRemoveList {
   std::atomic<FileToRemoveList *> Next = nullptr;
 
   FileToRemoveList() = default;
-  // Not signal-safe.
-  FileToRemoveList(const std::string &str) : Filename(strdup(str.c_str())) {}
+  // Takes ownership of \p filename.
+  FileToRemoveList(char *filename) : Filename(filename) {}
 
 public:
   // Not signal-safe.
@@ -116,19 +120,35 @@ public:
     if (FileToRemoveList *N = Next.exchange(nullptr))
       delete N;
     if (char *F = Filename.exchange(nullptr))
-      free(F);
+      if (F != InvalidPathSentinel)
+        free(F);
   }
 
   // Not signal-safe.
   static void insert(std::atomic<FileToRemoveList *> &Head,
                      const std::string &Filename) {
-    // Insert the new file at the end of the list.
-    FileToRemoveList *NewHead = new FileToRemoveList(Filename);
+    // Reuse a node with a null filename (left behind by erase) if one exists.
+    // There are two cases where Filename can be special:
+    // - nullptr: a node left behind by a previous file that we had to remove
+    // - InvalidPathSentinel: a node whose file is actively being removed by a
+    //   signal handler right now, in which case it's OK if this file doesn't
+    //   get removed.
+    char *NewFilename = strdup(Filename.c_str());
     std::atomic<FileToRemoveList *> *InsertionPoint = &Head;
-    FileToRemoveList *OldHead = nullptr;
-    while (!InsertionPoint->compare_exchange_strong(OldHead, NewHead)) {
-      InsertionPoint = &OldHead->Next;
-      OldHead = nullptr;
+    for (FileToRemoveList *Current = Head.load(); Current;
+         Current = Current->Next.load()) {
+      char *NullFilename = nullptr;
+      if (Current->Filename.compare_exchange_strong(NullFilename, NewFilename))
+        return; // Reused a slot.
+      InsertionPoint = &Current->Next;
+    }
+
+    // Append the new node at the end; on CAS failure, advance to the new tail.
+    FileToRemoveList *NewNode = new FileToRemoveList(NewFilename);
+    FileToRemoveList *OldNext = nullptr;
+    while (!InsertionPoint->compare_exchange_strong(OldNext, NewNode)) {
+      InsertionPoint = &OldNext->Next;
+      OldNext = nullptr;
     }
   }
 
@@ -145,11 +165,15 @@ public:
       if (char *OldFilename = Current->Filename.load()) {
         if (OldFilename != Filename)
           continue;
-        // Leave an empty filename.
-        OldFilename = Current->Filename.exchange(nullptr);
-        // The filename might have become null between the time we
-        // compared it and we exchanged it.
-        if (OldFilename)
+        // Leave an empty filename. Use CAS to avoid racing with the signal
+        // handler (which can't take the writer lock); only clear and free
+        // if we still own the pointer.
+        char *Expected = OldFilename;
+        while (!Current->Filename.compare_exchange_strong(Expected, nullptr)) {
+          if (Expected == nullptr || Expected == InvalidPathSentinel)
+            break;
+        }
+        if (Expected == OldFilename)
           free(OldFilename);
       }
     }
@@ -175,21 +199,23 @@ public:
 
   // Signal-safe.
   static void removeAllFiles(std::atomic<FileToRemoveList *> &Head) {
-    // If cleanup were to occur while we're removing files we'd have a bad time.
-    // Make sure we're OK by preventing cleanup from doing anything while we're
-    // removing files. If cleanup races with us and we win we'll have a leak,
-    // but we won't crash.
+    // This signal-safe code cannot acquire the writer lock, and needs to defend
+    // against racing writes from the `erase` method above.
     FileToRemoveList *OldHead = Head.exchange(nullptr);
 
     for (FileToRemoveList *currentFile = OldHead; currentFile;
          currentFile = currentFile->Next.load()) {
-      // If erasing was occuring while we're trying to remove files we'd look
-      // at free'd data. Take away the path and put it back when done.
-      if (char *path = currentFile->Filename.exchange(nullptr)) {
-        removeFile(path);
-
-        // We're done removing the file, erasing can safely proceed.
-        currentFile->Filename.exchange(path);
+      // Take exclusive ownership by swapping in the sentinel (signal-safe: no
+      // allocation or free). Then put the path back so we don't leak.
+      char *Path = currentFile->Filename.exchange(InvalidPathSentinel);
+      if (!Path) {
+        // Restore an empty slot so future insertions can reuse it.
+        currentFile->Filename.exchange(nullptr);
+      } else if (Path != InvalidPathSentinel) {
+        removeFile(Path);
+        // Add the path back to the list to create a global root referencing the
+        // heap allocation, which will pacify leak checkers that run at exit.
+        currentFile->Filename.exchange(Path);
       }
     }