[llvm-profgen] Remove temporary perf script files

[HaohaiWen]

The temporary perf script files converted from perf data will occupy lots
of space for large project. This patch removes them when llvm-profgen
exits normally or receives signals.

[llvmbot]

@llvm/pr-subscribers-llvm-support

@llvm/pr-subscribers-pgo

Author: Haohai Wen (HaohaiWen)

Changes

The temporary perf script files converted from perf data will occupy lots
of space for large project. This patch removes them when llvm-profgen
exits normally or receives signals.

---

Full diff: https://github.com/llvm/llvm-project/pull/86668.diff

3 Files Affected:

• (modified) llvm/tools/llvm-profgen/PerfReader.cpp (+19)
• (modified) llvm/tools/llvm-profgen/PerfReader.h (+7)
• (modified) llvm/tools/llvm-profgen/llvm-profgen.cpp (+2)

diff --git a/llvm/tools/llvm-profgen/PerfReader.cpp b/llvm/tools/llvm-profgen/PerfReader.cpp
index 878147642aa6e7..5ed0f1333cf2eb 100644
--- a/llvm/tools/llvm-profgen/PerfReader.cpp
+++ b/llvm/tools/llvm-profgen/PerfReader.cpp
@@ -11,6 +11,7 @@
 #include "llvm/DebugInfo/Symbolize/SymbolizableModule.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Process.h"
+#include "llvm/Support/Signals.h"
 
 #define DEBUG_TYPE "perf-reader"
 
@@ -375,6 +376,9 @@ PerfScriptReader::convertPerfDataToTrace(ProfiledBinary *Binary,
                                           StringRef(ErrorFile)};    // Stderr
   sys::ExecuteAndWait(PerfPath, ScriptMMapArgs, std::nullopt, Redirects);
 
+  PerfScriptReader::markTempFile(PerfTraceFile);
+  PerfScriptReader::markTempFile(ErrorFile);
+
   // Collect the PIDs
   TraceStream TraceIt(PerfTraceFile);
   std::string PIDs;
@@ -992,6 +996,11 @@ bool PerfScriptReader::extractMMap2EventForBinary(ProfiledBinary *Binary,
   return Binary->getName() == BinaryName;
 }
 
+void PerfScriptReader::markTempFile(StringRef FileName) {
+  sys::RemoveFileOnSignal(FileName);
+  TempFiles.push_back(FileName.str());
+}
+
 void PerfScriptReader::parseMMap2Event(TraceStream &TraceIt) {
   MMapEvent MMap;
   if (extractMMap2EventForBinary(Binary, TraceIt.getCurrentLine(), MMap))
@@ -1081,6 +1090,14 @@ PerfContent PerfScriptReader::checkPerfScriptType(StringRef FileName) {
   return PerfContent::UnknownContent;
 }
 
+void PerfScriptReader::removeTempFiles() {
+  for (StringRef FileName : TempFiles) {
+    sys::fs::remove(FileName);
+    sys::DontRemoveFileOnSignal(FileName);
+  }
+  TempFiles.clear();
+}
+
 void HybridPerfReader::generateUnsymbolizedProfile() {
   ProfileIsCS = !IgnoreStackSamples;
   if (ProfileIsCS)
@@ -1220,5 +1237,7 @@ void PerfScriptReader::parsePerfTraces() {
     writeUnsymbolizedProfile(OutputFilename);
 }
 
+SmallVector<std::string, 2> PerfScriptReader::TempFiles;
+
 } // end namespace sampleprof
 } // end namespace llvm
diff --git a/llvm/tools/llvm-profgen/PerfReader.h b/llvm/tools/llvm-profgen/PerfReader.h
index e9f619350bf970..6401a4c51c8b10 100644
--- a/llvm/tools/llvm-profgen/PerfReader.h
+++ b/llvm/tools/llvm-profgen/PerfReader.h
@@ -603,6 +603,8 @@ class PerfScriptReader : public PerfReaderBase {
                          std::optional<uint32_t> PIDFilter);
   // Extract perf script type by peaking at the input
   static PerfContent checkPerfScriptType(StringRef FileName);
+  // Remove all temporary files.
+  static void removeTempFiles();
 
 protected:
   // The parsed MMap event
@@ -622,6 +624,8 @@ class PerfScriptReader : public PerfReaderBase {
   // mapping between the binary name and its memory layout.
   static bool extractMMap2EventForBinary(ProfiledBinary *Binary, StringRef Line,
                                          MMapEvent &MMap);
+  // Mark temporary file for future clean up.
+  static void markTempFile(StringRef FileName);
   // Update base address based on mmap events
   void updateBinaryAddress(const MMapEvent &Event);
   // Parse mmap event and update binary address
@@ -662,6 +666,9 @@ class PerfScriptReader : public PerfReaderBase {
   std::set<uint64_t> InvalidReturnAddresses;
   // PID for the process of interest
   std::optional<uint32_t> PIDFilter;
+
+  // Temporary files created by perf script command.
+  static SmallVector<std::string, 2> TempFiles;
 };
 
 /*
diff --git a/llvm/tools/llvm-profgen/llvm-profgen.cpp b/llvm/tools/llvm-profgen/llvm-profgen.cpp
index 3b974e25103ad4..d0ec463f717677 100644
--- a/llvm/tools/llvm-profgen/llvm-profgen.cpp
+++ b/llvm/tools/llvm-profgen/llvm-profgen.cpp
@@ -189,5 +189,7 @@ int main(int argc, const char *argv[]) {
     Generator->write();
   }
 
+  PerfScriptReader::removeTempFiles();
+
   return EXIT_SUCCESS;
 }

[HaohaiWen]

Fix unresolved issue in #83489

[MatzeB]

Thanks! LGTM

[WenleiHe]

LGTM with a nit.

[WenleiHe]

nit: I would slightly prefer we do "createTempFile" instead of "createFile" + "MarkTemp". no strong opinion though.

[MatzeB]

I don't think sys::fs::createTemporaryFile has any of the cleanup logic from MarkTemp. There is the TempFile class but as far as I can tell that is meant towards the use case of getting a C++ file descriptor that you can write to but doesn't work to give you a filename + later cleanup; you only get a filename when you explicitely do not want the cleanup to happen.

So while I looked a bit when reviewing, I don't think there is much better than can be done here?

[WenleiHe]

I meant to suggest combing "createUniquePath" and "markTempFile" into a single helper, so people won't create a temp file and forget to mark it as temp..

[MaskRay]

Thx. This mechanism is similar to "llvm/Support/ToolOutputFile.h"

[wlei-llvm]

Nit: there are more "return EXIT_SUCCESS;" in other places, may be useful to also cover them(like use/call a wrapper function)? though other usages are not commonly used..

[HaohaiWen]

Using CleanupInstaller can avoid this since cleanup will be auto run in destructor of static variable.

[HaohaiWen]

Thx. This mechanism is similar to "llvm/Support/ToolOutputFile.h"

Thanks for letting me know. This make the code much cleaner. I filed a patch to make CleanupInstaller public so that it can be used here. #86758