[Support] Add overload writeFileAtomically(std::function Writer)

Differential Revision: https://reviews.llvm.org/D67424

llvm-svn: 371890

include/llvm/Support/FileUtilities.h

@@ -16,6 +16,7 @@
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Errc.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 
@@ -75,10 +76,40 @@ namespace llvm {
     void releaseFile() { DeleteIt = false; }
   };
 
+  enum class atomic_write_error {
+    failed_to_create_uniq_file = 0,
+    output_stream_error,
+    failed_to_rename_temp_file
+  };
+
+  class AtomicFileWriteError : public llvm::ErrorInfo<AtomicFileWriteError> {
+  public:
+    AtomicFileWriteError(atomic_write_error Error) : Error(Error) {}
+
+    void log(raw_ostream &OS) const override;
+
+    const atomic_write_error Error;
+    static char ID;
+
+  private:
+    // Users are not expected to use error_code.
+    std::error_code convertToErrorCode() const override {
+      return llvm::inconvertibleErrorCode();
+    }
+  };
+
+  // atomic_write_error + whatever the Writer can return
+
   /// Creates a unique file with name according to the given \p TempPathModel,
   /// writes content of \p Buffer to the file and renames it to \p FinalPath.
+  ///
+  /// \returns \c AtomicFileWriteError in case of error.
   llvm::Error writeFileAtomically(StringRef TempPathModel, StringRef FinalPath,
                                   StringRef Buffer);
+
+  llvm::Error
+  writeFileAtomically(StringRef TempPathModel, StringRef FinalPath,
+                      std::function<llvm::Error(llvm::raw_ostream &)> Writer);
 } // End llvm namespace
 
 #endif

lib/LTO/ThinLTOCodeGenerator.cpp

@@ -39,6 +39,7 @@
 #include "llvm/Support/CachePruning.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/SHA1.h"
 #include "llvm/Support/SmallVectorMemoryBuffer.h"
@@ -368,23 +369,26 @@ public:
     // Write to a temporary to avoid race condition
     SmallString<128> TempFilename;
     SmallString<128> CachePath(EntryPath);
-    int TempFD;
     llvm::sys::path::remove_filename(CachePath);
     sys::path::append(TempFilename, CachePath, "Thin-%%%%%%.tmp.o");
-    std::error_code EC =
-      sys::fs::createUniqueFile(TempFilename, TempFD, TempFilename);
-    if (EC) {
-      errs() << "Error: " << EC.message() << "\n";
+
+    if (auto Err = handleErrors(
+            llvm::writeFileAtomically(TempFilename, EntryPath,
+                                      OutputBuffer.getBuffer()),
+            [](const llvm::AtomicFileWriteError &E) {
+              std::string ErrorMsgBuffer;
+              llvm::raw_string_ostream S(ErrorMsgBuffer);
+              E.log(S);
+
+              if (E.Error ==
+                  llvm::atomic_write_error::failed_to_create_uniq_file) {
+                errs() << "Error: " << ErrorMsgBuffer << "\n";
                 report_fatal_error("ThinLTO: Can't get a temporary file");
               }
-    {
-      raw_fd_ostream OS(TempFD, /* ShouldClose */ true);
-      OS << OutputBuffer.getBuffer();
+            })) {
+      // FIXME
+      consumeError(std::move(Err));
     }
-    // Rename temp file to final destination; rename is atomic
-    EC = sys::fs::rename(TempFilename, EntryPath);
-    if (EC)
-      sys::fs::remove(TempFilename);
   }
 };
 

lib/Support/FileUtilities.cpp

@@ -12,7 +12,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/FileUtilities.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
@@ -266,36 +268,65 @@ int llvm::DiffFilesWithTolerance(StringRef NameA,
   return CompareFailed;
 }
 
-Error llvm::writeFileAtomically(StringRef TempPathModel, StringRef FinalPath,
-                                StringRef Buffer) {
+void llvm::AtomicFileWriteError::log(raw_ostream &OS) const {
+  OS << "atomic_write_error: ";
+  switch (Error) {
+  case atomic_write_error::failed_to_create_uniq_file:
+    OS << "failed_to_create_uniq_file";
+    return;
+  case atomic_write_error::output_stream_error:
+    OS << "output_stream_error";
+    return;
+  case atomic_write_error::failed_to_rename_temp_file:
+    OS << "failed_to_rename_temp_file";
+    return;
+  }
+  llvm_unreachable("unknown atomic_write_error value in "
+                   "failed_to_rename_temp_file::log()");
+}
+
+llvm::Error llvm::writeFileAtomically(StringRef TempPathModel,
+                                      StringRef FinalPath, StringRef Buffer) {
+  return writeFileAtomically(TempPathModel, FinalPath,
+                             [&Buffer](llvm::raw_ostream &OS) {
+                               OS.write(Buffer.data(), Buffer.size());
+                               return llvm::Error::success();
+                             });
+}
+
+llvm::Error llvm::writeFileAtomically(
+    StringRef TempPathModel, StringRef FinalPath,
+    std::function<llvm::Error(llvm::raw_ostream &)> Writer) {
   SmallString<128> GeneratedUniqPath;
   int TempFD;
-  if (const std::error_code Error = sys::fs::createUniqueFile(
-          TempPathModel.str(), TempFD, GeneratedUniqPath)) {
-    return createStringError(
-        Error, "failed to create temporary file with model \"%s\"",
-        TempPathModel.str().c_str());
+  if (sys::fs::createUniqueFile(TempPathModel.str(), TempFD,
+                                GeneratedUniqPath)) {
+    return llvm::make_error<AtomicFileWriteError>(
+        atomic_write_error::failed_to_create_uniq_file);
   }
+  llvm::FileRemover RemoveTmpFileOnFail(GeneratedUniqPath);
 
   raw_fd_ostream OS(TempFD, /*shouldClose=*/true);
-  OS.write(Buffer.data(), Buffer.size());
-  OS.close();
-  TempFD = -1;
+  if (llvm::Error Err = Writer(OS)) {
+    return Err;
+  }
 
+  OS.close();
   if (OS.has_error()) {
-    const std::error_code Error = OS.error();
     OS.clear_error();
-    return createStringError(Error, "failed to write to \"%s\"",
-                             GeneratedUniqPath.c_str());
+    return llvm::make_error<AtomicFileWriteError>(
+        atomic_write_error::output_stream_error);
   }
 
   if (const std::error_code Error =
           sys::fs::rename(/*from=*/GeneratedUniqPath.c_str(),
                           /*to=*/FinalPath.str().c_str())) {
-    return createStringError(Error, "failed to rename file \"%s\" to \"%s\"",
-                             GeneratedUniqPath.c_str(),
-                             FinalPath.str().c_str());
+    return llvm::make_error<AtomicFileWriteError>(
+        atomic_write_error::failed_to_rename_temp_file);
   }
 
+  RemoveTmpFileOnFail.releaseFile();
   return Error::success();
 }
+
+char llvm::AtomicFileWriteError::ID;