[DTLTO] Overlap temporary file removal (#209423)

Deleting the temporary files produced by the DTLTO pipeline can be
expensive on Windows hosts. For a Clang link (Debug build with
sanitizers and instrumentation) using an optimized toolchain (PGO
non-LTO, llvmorg-22.1.0) on a Windows 11 Pro (Build 26200), AMD Family
25 @ ~4.5 GHz, 16 cores/32 threads, 64 GB RAM machine, the mean duration
of the "Remove DTLTO temporary files" time trace scope was 1267.789 ms
(measured over 10 runs).

This patch performs the deletions on a background thread, allowing them
to overlap with the tail of the link to hide this cost.

This is a re-implementation of the asynchronous cleanup idea from
https://github.com/llvm/llvm-project/pull/186988, which had to be
reverted in https://github.com/llvm/llvm-project/pull/189043 because
cleanup was not guaranteed to complete before LLD invoked
timeTraceProfilerCleanup(). In certain cases timeTraceProfilerCleanup()
was called before temporary file deletion had completed in LLD, which
caused memory leaks that were flagged by sanitizer builds.

Note that the DTLTO implementation has been refactored heavily since
https://github.com/llvm/llvm-project/pull/186988, so this is a
re-implementation along the same lines.

To solve the ordering issue, LLD now calls a hook that defaults to a
no-op, and DTLTO overrides it to drain the background deletion work.
LLD calls this hook before time-trace write/cleanup.

lld/COFF/Driver.cpp

@@ -2974,6 +2974,9 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
 
   // Write the result.
   writeResult(ctx);
+  // LTO cleanup may create time trace events. Wait for it to complete before
+  // writing the time trace data.
+  ctx.forEachSymtab([](SymbolTable &symtab) { symtab.waitForLTOCleanup(); });
 
   // Stop early so we can print the results.
   rootTimer.stop();

lld/COFF/LTO.cpp

@@ -156,6 +156,11 @@ BitcodeCompiler::BitcodeCompiler(COFFLinkerContext &c) : ctx(c) {
 
 BitcodeCompiler::~BitcodeCompiler() = default;
 
+void BitcodeCompiler::waitForLTOCleanup() {
+  if (ltoObj)
+    ltoObj->waitForCleanup();
+}
+
 static void undefine(Symbol *s) { replaceSymbol<Undefined>(s, s->getName()); }
 
 void BitcodeCompiler::add(BitcodeFile &f) {

lld/COFF/LTO.h

@@ -46,6 +46,7 @@ public:
   void add(BitcodeFile &f);
   std::vector<InputFile *> compile();
   void setBitcodeLibFuncs(ArrayRef<StringRef> bitcodeLibFuncs);
+  void waitForLTOCleanup();
 
 private:
   std::unique_ptr<llvm::lto::LTO> ltoObj;

lld/COFF/SymbolTable.cpp

@@ -1511,4 +1511,9 @@ void SymbolTable::compileBitcodeFiles() {
   }
 }
 
+void SymbolTable::waitForLTOCleanup() {
+  if (lto)
+    lto->waitForLTOCleanup();
+}
+
 } // namespace lld::coff

lld/COFF/SymbolTable.h

@@ -109,6 +109,8 @@ public:
   // added and before the writer writes results to a file.
   void compileBitcodeFiles();
 
+  void waitForLTOCleanup();
+
   // Creates an Undefined symbol and marks it as live.
   Symbol *addGCRoot(StringRef sym, bool aliasEC = false);
 

lld/ELF/Config.h

@@ -215,6 +215,7 @@ private:
   void createFiles(llvm::opt::InputArgList &args);
   void loadFiles();
   void inferMachineType();
+  void waitForLTOCleanup();
   template <class ELFT> void link(llvm::opt::InputArgList &args);
   template <class ELFT> void compileBitcodeFiles(bool skipLinkedOutput);
   // True if we are in --whole-archive and --no-whole-archive.

lld/ELF/Driver.cpp

@@ -625,6 +625,11 @@ constexpr const char *saveTempsValues[] = {
 
 LinkerDriver::LinkerDriver(Ctx &ctx) : ctx(ctx) {}
 
+void LinkerDriver::waitForLTOCleanup() {
+  if (lto)
+    lto->waitForLTOCleanup();
+}
+
 void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   ELFOptTable parser;
   opt::InputArgList args = parser.parse(ctx, argsArr.slice(1));
@@ -707,6 +712,10 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
     invokeELFT(link, args);
   }
 
+  // LTO cleanup may create time trace events. Wait for it to complete before
+  // writing the time trace data.
+  waitForLTOCleanup();
+
   if (ctx.arg.timeTraceEnabled) {
     checkError(ctx.e, timeTraceProfilerWrite(
                           args.getLastArgValue(OPT_time_trace_eq).str(),

lld/ELF/LTO.cpp

@@ -231,6 +231,11 @@ BitcodeCompiler::BitcodeCompiler(Ctx &ctx) : ctx(ctx) {
 
 BitcodeCompiler::~BitcodeCompiler() = default;
 
+void BitcodeCompiler::waitForLTOCleanup() {
+  if (ltoObj)
+    ltoObj->waitForCleanup();
+}
+
 void BitcodeCompiler::add(BitcodeFile &f) {
   lto::InputFile &obj = *f.obj;
   bool isExec = !ctx.arg.shared && !ctx.arg.relocatable;

lld/ELF/LTO.h

@@ -44,6 +44,7 @@ public:
   void add(BitcodeFile &f);
   SmallVector<std::unique_ptr<InputFile>, 0> compile();
   void setBitcodeLibFuncs(ArrayRef<StringRef> bitcodeLibFuncs);
+  void waitForLTOCleanup();
 
 private:
   Ctx &ctx;

llvm/include/llvm/DTLTO/DTLTO.h

@@ -19,10 +19,12 @@
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/LTO/LTO.h"
-#include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/Signals.h"
+#include "llvm/Support/ThreadPool.h"
 
 #include <functional>
+#include <string>
+#include <utility>
 #include <vector>
 
 namespace llvm {
@@ -94,6 +96,12 @@ public:
   /// Cache) for each task identifier.
   virtual Error run(AddStreamFn AddStream, FileCache Cache = {}) override;
 
+  /// Wait for LTO cleanup. Clients may call this after run() once subsequent
+  /// linking work that can overlap with cleanup is complete. Cleanup may emit
+  /// time trace events, so this must be called before time trace data is
+  /// finalized.
+  void waitForCleanup() override;
+
 private:
   /// DTLTO archive support.
   ///
@@ -356,6 +364,20 @@ private:
   // Cleanup files list.
   std::vector<std::string> CleanupList;
 
+  // There can be many temporary files to remove. Performing deletion in the
+  // background can save a few seconds on Windows hosts.
+  struct BackgroundDeletion : DefaultThreadPool {
+    BackgroundDeletion();
+    ~BackgroundDeletion();
+
+    void removeFiles(std::vector<std::string> &&Files, const Config &Conf);
+    void waitForTasks();
+
+    std::vector<std::string> Warnings;
+  };
+
+  BackgroundDeletion BackgroundDeleter;
+
   // Record a file for cleanup and register signal-time removal if requested.
   void addToCleanup(StringRef Filename) {
     CleanupList.push_back(Filename.str());

llvm/include/llvm/LTO/LTO.h

@@ -439,6 +439,14 @@ public:
   /// Cache) for each task identifier.
   virtual Error run(AddStreamFn AddStream, FileCache Cache = {});
 
+  /// Wait for cleanup work started by run() to finish.
+  ///
+  /// A client may delay this call to overlap asynchronous cleanup with later
+  /// linking work, but must call it before finalizing time trace data because
+  /// cleanup may emit time trace events. Most LTO implementations have no
+  /// asynchronous cleanup.
+  virtual void waitForCleanup() {}
+
   /// Static method that returns a list of libcall symbols that can be generated
   /// by LTO but might not be visible from bitcode symbol table.
   static SmallVector<const char *> getRuntimeLibcallSymbols(const Triple &TT);

llvm/lib/DTLTO/DTLTO.cpp

@@ -27,27 +27,59 @@
 #include "llvm/Support/raw_ostream.h"
 
 #include <string>
+#include <system_error>
+#include <utility>
+#include <vector>
 
 using namespace llvm;
 
+// Experimentation showed that serial deletion is most efficient, hence
+// a single thread.
+lto::DTLTO::BackgroundDeletion::BackgroundDeletion()
+    : DefaultThreadPool(hardware_concurrency(1)) {}
+
+lto::DTLTO::BackgroundDeletion::~BackgroundDeletion() { waitForTasks(); }
+
+void lto::DTLTO::BackgroundDeletion::waitForTasks() {
+  wait();
+  for (const std::string &Warning : Warnings)
+    errs() << "warning: could not remove the file " << Warning << "\n";
+  Warnings.clear();
+}
+
+void lto::DTLTO::BackgroundDeletion::removeFiles(
+    std::vector<std::string> &&Files, const Config &Conf) {
+  if (Files.empty())
+    return;
+
+  async([this, Files = std::move(Files), TTE = Conf.TimeTraceEnabled,
+         TTG = Conf.TimeTraceGranularity] {
+    if (LLVM_ENABLE_THREADS && TTE)
+      timeTraceProfilerInitialize(TTG, "Remove DTLTO temporary files");
+    {
+      TimeTraceScope TimeScope("Remove DTLTO temporary files");
+      for (const std::string &Path : Files) {
+        std::error_code EC = sys::fs::remove(Path, true);
+        if (!EC ||
+            EC == std::make_error_code(std::errc::no_such_file_or_directory))
+          continue;
+
+        Warnings.emplace_back("'" + Path + "': " + EC.message());
+      }
+    }
+    if (LLVM_ENABLE_THREADS && TTE)
+      timeTraceProfilerFinishThread();
+  });
+}
+
+void lto::DTLTO::waitForCleanup() { BackgroundDeleter.waitForTasks(); }
+
 // Remove temporary files created to enable distribution.
 void lto::DTLTO::cleanup() {
-  if (!SaveTemps) {
-    // Remove one file, report error if any.
-    auto removeFile = [](StringRef FileName) -> void {
-      std::error_code EC = sys::fs::remove(FileName, true);
-      if (EC &&
-          EC != std::make_error_code(std::errc::no_such_file_or_directory))
-        errs() << "warning: could not remove the file '" << FileName
-               << "': " << EC.message() << "\n";
-    };
+  if (SaveTemps)
+    return;
 
-    TimeTraceScope JobScope("Remove DTLTO temporary files");
-    for (const auto &Name : CleanupList)
-      removeFile(Name);
-    // Clean the CleanupList for safety.
-    CleanupList.clear();
-  }
+  BackgroundDeleter.removeFiles(std::move(CleanupList), Conf);
 }
 
 // Runs the DTLTO thin link phase, producing per-module summary indices,

llvm/tools/llvm-lto2/llvm-lto2.cpp

@@ -541,6 +541,7 @@ static int run(int argc, char **argv) {
                   "failed to create cache");
 
   check(Lto->run(AddStream, Cache), "LTO::run failed");
+  Lto->waitForCleanup();
   return static_cast<int>(HasErrors);
 }