15fce8c4fSNick Kledzik //===- FileOutputBuffer.cpp - File Output Buffer ----------------*- C++ -*-===//
25fce8c4fSNick Kledzik //
35fce8c4fSNick Kledzik //                     The LLVM Compiler Infrastructure
45fce8c4fSNick Kledzik //
55fce8c4fSNick Kledzik // This file is distributed under the University of Illinois Open Source
65fce8c4fSNick Kledzik // License. See LICENSE.TXT for details.
75fce8c4fSNick Kledzik //
85fce8c4fSNick Kledzik //===----------------------------------------------------------------------===//
95fce8c4fSNick Kledzik //
105fce8c4fSNick Kledzik // Utility for creating a in-memory buffer that will be written to a file.
115fce8c4fSNick Kledzik //
125fce8c4fSNick Kledzik //===----------------------------------------------------------------------===//
135fce8c4fSNick Kledzik 
14d9903888SChandler Carruth #include "llvm/Support/FileOutputBuffer.h"
1516132e6fSBenjamin Kramer #include "llvm/ADT/STLExtras.h"
1616132e6fSBenjamin Kramer #include "llvm/ADT/SmallString.h"
1716132e6fSBenjamin Kramer #include "llvm/Support/Errc.h"
18a16fe65bSRui Ueyama #include "llvm/Support/Memory.h"
19d4b24edaSRafael Espindola #include "llvm/Support/Path.h"
207dbb5778SRafael Espindola #include "llvm/Support/Signals.h"
21a6e9c3e4SRafael Espindola #include <system_error>
225fce8c4fSNick Kledzik 
237eb1f185SRafael Espindola #if !defined(_MSC_VER) && !defined(__MINGW32__)
247eb1f185SRafael Espindola #include <unistd.h>
257eb1f185SRafael Espindola #else
267eb1f185SRafael Espindola #include <io.h>
277eb1f185SRafael Espindola #endif
287eb1f185SRafael Espindola 
29a16fe65bSRui Ueyama using namespace llvm;
30a16fe65bSRui Ueyama using namespace llvm::sys;
315fce8c4fSNick Kledzik 
32a16fe65bSRui Ueyama // A FileOutputBuffer which creates a temporary file in the same directory
33a16fe65bSRui Ueyama // as the final output file. The final output file is atomically replaced
34a16fe65bSRui Ueyama // with the temporary file on commit().
35a16fe65bSRui Ueyama class OnDiskBuffer : public FileOutputBuffer {
36a16fe65bSRui Ueyama public:
37a16fe65bSRui Ueyama   OnDiskBuffer(StringRef Path, StringRef TempPath,
38a16fe65bSRui Ueyama                std::unique_ptr<fs::mapped_file_region> Buf)
39a16fe65bSRui Ueyama       : FileOutputBuffer(Path), Buffer(std::move(Buf)), TempPath(TempPath) {}
405fce8c4fSNick Kledzik 
41*e0df357dSRafael Espindola   static Expected<std::unique_ptr<OnDiskBuffer>>
42a16fe65bSRui Ueyama   create(StringRef Path, size_t Size, unsigned Mode);
43a16fe65bSRui Ueyama 
44a16fe65bSRui Ueyama   uint8_t *getBufferStart() const override { return (uint8_t *)Buffer->data(); }
45a16fe65bSRui Ueyama 
46a16fe65bSRui Ueyama   uint8_t *getBufferEnd() const override {
47a16fe65bSRui Ueyama     return (uint8_t *)Buffer->data() + Buffer->size();
48a16fe65bSRui Ueyama   }
49a16fe65bSRui Ueyama 
50a16fe65bSRui Ueyama   size_t getBufferSize() const override { return Buffer->size(); }
51a16fe65bSRui Ueyama 
52a16fe65bSRui Ueyama   std::error_code commit() override {
53a16fe65bSRui Ueyama     // Unmap buffer, letting OS flush dirty pages to file on disk.
54a16fe65bSRui Ueyama     Buffer.reset();
55a16fe65bSRui Ueyama 
56a16fe65bSRui Ueyama     // Atomically replace the existing file with the new one.
57a16fe65bSRui Ueyama     auto EC = fs::rename(TempPath, FinalPath);
58a16fe65bSRui Ueyama     sys::DontRemoveFileOnSignal(TempPath);
59a16fe65bSRui Ueyama     return EC;
60a16fe65bSRui Ueyama   }
61a16fe65bSRui Ueyama 
62a16fe65bSRui Ueyama   ~OnDiskBuffer() override {
631a4398a1SReid Kleckner     // Close the mapping before deleting the temp file, so that the removal
641a4398a1SReid Kleckner     // succeeds.
65a16fe65bSRui Ueyama     Buffer.reset();
66a16fe65bSRui Ueyama     fs::remove(TempPath);
675fce8c4fSNick Kledzik   }
685fce8c4fSNick Kledzik 
69a16fe65bSRui Ueyama private:
70a16fe65bSRui Ueyama   std::unique_ptr<fs::mapped_file_region> Buffer;
71a16fe65bSRui Ueyama   std::string TempPath;
72a16fe65bSRui Ueyama };
73a16fe65bSRui Ueyama 
74a16fe65bSRui Ueyama // A FileOutputBuffer which keeps data in memory and writes to the final
75a16fe65bSRui Ueyama // output file on commit(). This is used only when we cannot use OnDiskBuffer.
76a16fe65bSRui Ueyama class InMemoryBuffer : public FileOutputBuffer {
77a16fe65bSRui Ueyama public:
78a16fe65bSRui Ueyama   InMemoryBuffer(StringRef Path, MemoryBlock Buf, unsigned Mode)
79a16fe65bSRui Ueyama       : FileOutputBuffer(Path), Buffer(Buf), Mode(Mode) {}
80a16fe65bSRui Ueyama 
81*e0df357dSRafael Espindola   static Expected<std::unique_ptr<InMemoryBuffer>>
82a16fe65bSRui Ueyama   create(StringRef Path, size_t Size, unsigned Mode) {
83a16fe65bSRui Ueyama     std::error_code EC;
84a16fe65bSRui Ueyama     MemoryBlock MB = Memory::allocateMappedMemory(
85a16fe65bSRui Ueyama         Size, nullptr, sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC);
865fce8c4fSNick Kledzik     if (EC)
87*e0df357dSRafael Espindola       return errorCodeToError(EC);
88a16fe65bSRui Ueyama     return llvm::make_unique<InMemoryBuffer>(Path, MB, Mode);
895fce8c4fSNick Kledzik   }
905fce8c4fSNick Kledzik 
91a16fe65bSRui Ueyama   uint8_t *getBufferStart() const override { return (uint8_t *)Buffer.base(); }
92a16fe65bSRui Ueyama 
93a16fe65bSRui Ueyama   uint8_t *getBufferEnd() const override {
94a16fe65bSRui Ueyama     return (uint8_t *)Buffer.base() + Buffer.size();
95a16fe65bSRui Ueyama   }
96a16fe65bSRui Ueyama 
97a16fe65bSRui Ueyama   size_t getBufferSize() const override { return Buffer.size(); }
98a16fe65bSRui Ueyama 
99a16fe65bSRui Ueyama   std::error_code commit() override {
100d4b24edaSRafael Espindola     int FD;
101a16fe65bSRui Ueyama     std::error_code EC;
102a16fe65bSRui Ueyama     if (auto EC = openFileForWrite(FinalPath, FD, fs::F_None, Mode))
103a16fe65bSRui Ueyama       return EC;
104a16fe65bSRui Ueyama     raw_fd_ostream OS(FD, /*shouldClose=*/true, /*unbuffered=*/true);
105a16fe65bSRui Ueyama     OS << StringRef((const char *)Buffer.base(), Buffer.size());
106a16fe65bSRui Ueyama     return std::error_code();
107d4b24edaSRafael Espindola   }
108d4b24edaSRafael Espindola 
109a16fe65bSRui Ueyama private:
110a16fe65bSRui Ueyama   OwningMemoryBlock Buffer;
111a16fe65bSRui Ueyama   unsigned Mode;
112a16fe65bSRui Ueyama };
113a16fe65bSRui Ueyama 
114*e0df357dSRafael Espindola Expected<std::unique_ptr<OnDiskBuffer>>
115a16fe65bSRui Ueyama OnDiskBuffer::create(StringRef Path, size_t Size, unsigned Mode) {
116a16fe65bSRui Ueyama   // Create new file in same directory but with random name.
117a16fe65bSRui Ueyama   SmallString<128> TempPath;
118a16fe65bSRui Ueyama   int FD;
119a16fe65bSRui Ueyama   if (auto EC = fs::createUniqueFile(Path + ".tmp%%%%%%%", FD, TempPath, Mode))
120*e0df357dSRafael Espindola     return errorCodeToError(EC);
1215fce8c4fSNick Kledzik 
122a16fe65bSRui Ueyama   sys::RemoveFileOnSignal(TempPath);
1237dbb5778SRafael Espindola 
124da9bc2e5SRui Ueyama #ifndef LLVM_ON_WIN32
125da9bc2e5SRui Ueyama   // On Windows, CreateFileMapping (the mmap function on Windows)
126da9bc2e5SRui Ueyama   // automatically extends the underlying file. We don't need to
127da9bc2e5SRui Ueyama   // extend the file beforehand. _chsize (ftruncate on Windows) is
128da9bc2e5SRui Ueyama   // pretty slow just like it writes specified amount of bytes,
129a16fe65bSRui Ueyama   // so we should avoid calling that function.
130a16fe65bSRui Ueyama   if (auto EC = fs::resize_file(FD, Size))
131*e0df357dSRafael Espindola     return errorCodeToError(EC);
132da9bc2e5SRui Ueyama #endif
133c69f13bfSRafael Espindola 
134a16fe65bSRui Ueyama   // Mmap it.
135a16fe65bSRui Ueyama   std::error_code EC;
136a16fe65bSRui Ueyama   auto MappedFile = llvm::make_unique<fs::mapped_file_region>(
137a16fe65bSRui Ueyama       FD, fs::mapped_file_region::readwrite, Size, 0, EC);
138a16fe65bSRui Ueyama   close(FD);
1395fce8c4fSNick Kledzik   if (EC)
140*e0df357dSRafael Espindola     return errorCodeToError(EC);
141a16fe65bSRui Ueyama   return llvm::make_unique<OnDiskBuffer>(Path, TempPath, std::move(MappedFile));
1425fce8c4fSNick Kledzik }
1435fce8c4fSNick Kledzik 
144a16fe65bSRui Ueyama // Create an instance of FileOutputBuffer.
145*e0df357dSRafael Espindola Expected<std::unique_ptr<FileOutputBuffer>>
146a16fe65bSRui Ueyama FileOutputBuffer::create(StringRef Path, size_t Size, unsigned Flags) {
147a16fe65bSRui Ueyama   unsigned Mode = fs::all_read | fs::all_write;
148a16fe65bSRui Ueyama   if (Flags & F_executable)
149a16fe65bSRui Ueyama     Mode |= fs::all_exe;
1505fce8c4fSNick Kledzik 
151a16fe65bSRui Ueyama   fs::file_status Stat;
152a16fe65bSRui Ueyama   fs::status(Path, Stat);
153d4b24edaSRafael Espindola 
154a16fe65bSRui Ueyama   // Usually, we want to create OnDiskBuffer to create a temporary file in
155a16fe65bSRui Ueyama   // the same directory as the destination file and atomically replaces it
156a16fe65bSRui Ueyama   // by rename(2).
157a16fe65bSRui Ueyama   //
158a16fe65bSRui Ueyama   // However, if the destination file is a special file, we don't want to
159a16fe65bSRui Ueyama   // use rename (e.g. we don't want to replace /dev/null with a regular
160a16fe65bSRui Ueyama   // file.) If that's the case, we create an in-memory buffer, open the
161a16fe65bSRui Ueyama   // destination file and write to it on commit().
162a16fe65bSRui Ueyama   switch (Stat.type()) {
163a16fe65bSRui Ueyama   case fs::file_type::directory_file:
164*e0df357dSRafael Espindola     return errorCodeToError(errc::is_a_directory);
165a16fe65bSRui Ueyama   case fs::file_type::regular_file:
166a16fe65bSRui Ueyama   case fs::file_type::file_not_found:
167a16fe65bSRui Ueyama   case fs::file_type::status_error:
168a16fe65bSRui Ueyama     return OnDiskBuffer::create(Path, Size, Mode);
169a16fe65bSRui Ueyama   default:
170a16fe65bSRui Ueyama     return InMemoryBuffer::create(Path, Size, Mode);
1715fce8c4fSNick Kledzik   }
172a16fe65bSRui Ueyama }
173