ObjCMessageTracing: append trace file content atomically to a canonical file location provided by the build environment

To simplify post-processing of these collected ObjC message traces, this PR teaches the compiler to atomically append the content of
ObjC message trace to a file location provided by the build environment, instead of emitting multiple files.

Author: nkcsgexi

lib/FrontendTool/LoadedModuleTrace.cpp

@@ -853,16 +853,16 @@ class ObjcMethodReferenceCollector: public SourceEntityWalker {
   }
 };
 
-static std::optional<int> createObjCMessageTraceFile(const InputFile &input,
-                                                     ModuleDecl *MD,
-                                        std::vector<SourceFile*> &filesToWalk) {
+static void createObjCMessageTraceFile(const InputFile &input, ModuleDecl *MD) {
   if (input.getLoadedModuleTracePath().empty()) {
     // we basically rely on the passing down of module trace file path
     // as an indicator that this job needs to emit an ObjC message trace file.
     // FIXME: add a separate swift-frontend flag for ObjC message trace path
     // specifically.
-    return {};
+    return;
   }
+  auto &ctx = MD->getASTContext();
+  std::vector<SourceFile*> filesToWalk;
   for (auto *FU : MD->getFiles()) {
     if (auto *SF = dyn_cast<SourceFile>(FU)) {
       switch (SF->Kind) {
@@ -880,29 +880,47 @@ static std::optional<int> createObjCMessageTraceFile(const InputFile &input,
   }
   // No source files to walk, abort.
   if (filesToWalk.empty()) {
-    return {};
+    return;
   }
   llvm::SmallString<128> tracePath;
-  if (const char *P = ::getenv("SWIFT_COMPILER_OBJC_MESSAGE_TRACE_DIRECTORY")) {
-    StringRef DirPath = P;
-    llvm::sys::path::append(tracePath, DirPath);
+  if (const char *P = ::getenv("SWIFT_COMPILER_OBJC_MESSAGE_TRACE_PATH")) {
+    StringRef FilePath = P;
+    llvm::sys::path::append(tracePath, FilePath);
   } else {
     llvm::sys::path::append(tracePath, input.getLoadedModuleTracePath());
     llvm::sys::path::remove_filename(tracePath);
-    llvm::sys::path::append(tracePath, ".SWIFT_FINE_DEPENDENCY_TRACE");
+    llvm::sys::path::append(tracePath, ".SWIFT_FINE_DEPENDENCY_TRACE.json");
+  }
+  // Write output via atomic append.
+  llvm::vfs::OutputConfig config;
+  config.setAppend().setAtomicWrite();
+  auto outputFile = ctx.getOutputBackend().createFile(tracePath, config);
+  if (!outputFile) {
+    ctx.Diags.diagnose(SourceLoc(), diag::error_opening_output, tracePath,
+                       toString(outputFile.takeError()));
+    return;
   }
-  if (!llvm::sys::fs::exists(tracePath)) {
-    if (llvm::sys::fs::create_directory(tracePath))
-      return {};
+  ObjcMethodReferenceCollector collector(MD);
+  for (auto *SF : filesToWalk) {
+    collector.setFileBeforeVisiting(SF);
+    collector.walk(*SF);
+  }
+
+  // print this json line.
+  std::string stringBuffer;
+  {
+    llvm::raw_string_ostream memoryBuffer(stringBuffer);
+    collector.serializeAsJson(memoryBuffer);
   }
-  SmallString<32> fileName(MD->getNameStr());
-  fileName.append("-%%%%-%%%%-%%%%.json");
-  llvm::sys::path::append(tracePath, fileName);
-  int tmpFD;
-  if (llvm::sys::fs::createUniqueFile(tracePath.str(), tmpFD, tracePath)) {
-    return {};
+  stringBuffer += "\n";
+
+  // Write output via atomic append.
+  *outputFile << stringBuffer;
+  if (auto err = outputFile->keep()) {
+    ctx.Diags.diagnose(SourceLoc(), diag::error_opening_output,
+                       tracePath, toString(std::move(err)));
+    return;
   }
-  return tmpFD;
 }
 
 bool swift::emitObjCMessageSendTraceIfNeeded(ModuleDecl *mainModule,
@@ -912,18 +930,7 @@ bool swift::emitObjCMessageSendTraceIfNeeded(ModuleDecl *mainModule,
          "We should've already exited earlier if there was an error.");
 
   opts.InputsAndOutputs.forEachInput([&](const InputFile &input) {
-    std::vector<SourceFile*> filesToWalk;
-    auto tmpFD = createObjCMessageTraceFile(input, mainModule, filesToWalk);
-    if (!tmpFD)
-      return false;
-    // Write the contents of the buffer.
-    llvm::raw_fd_ostream out(*tmpFD, /*shouldClose=*/true);
-    ObjcMethodReferenceCollector collector(mainModule);
-    for (auto *SF : filesToWalk) {
-      collector.setFileBeforeVisiting(SF);
-      collector.walk(*SF);
-    }
-    collector.serializeAsJson(out);
+    createObjCMessageTraceFile(input, mainModule);
     return true;
   });
   return false;

test/IDE/objc_send_collector_1.swift

@@ -2,10 +2,10 @@
 // RUN: %empty-directory(%t/CUSTOM_DIR)
 
 // RUN: %target-swift-frontend  -I %t/lib/swift -typecheck %s %S/Inputs/objc_send_collector_2.swift -module-name main -swift-version 5 -F %S/Inputs/mock-sdk -emit-loaded-module-trace-path %t/.MODULE_TRACE
-// RUN: cat %t/.SWIFT_FINE_DEPENDENCY_TRACE/* | %FileCheck %s
+// RUN: cat %t/.SWIFT_FINE_DEPENDENCY_TRACE.json | %FileCheck %s
 
-// RUN: SWIFT_COMPILER_OBJC_MESSAGE_TRACE_DIRECTORY=%t/CUSTOM_DIR %target-swift-frontend  -I %t/lib/swift -typecheck %s %S/Inputs/objc_send_collector_2.swift -module-name main -swift-version 5 -F %S/Inputs/mock-sdk -emit-loaded-module-trace-path %t/.MODULE_TRACE
-// RUN: cat %t/CUSTOM_DIR/* | %FileCheck %s
+// RUN: SWIFT_COMPILER_OBJC_MESSAGE_TRACE_PATH=%t/given_trace.json %target-swift-frontend  -I %t/lib/swift -typecheck %s %S/Inputs/objc_send_collector_2.swift -module-name main -swift-version 5 -F %S/Inputs/mock-sdk -emit-loaded-module-trace-path %t/.MODULE_TRACE
+// RUN: cat %t/given_trace.json | %FileCheck %s
 
 // REQUIRES: objc_interop