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 
41a16fe65bSRui Ueyama   uint8_t *getBufferStart() const override { return (uint8_t *)Buffer->data(); }
42a16fe65bSRui Ueyama 
43a16fe65bSRui Ueyama   uint8_t *getBufferEnd() const override {
44a16fe65bSRui Ueyama     return (uint8_t *)Buffer->data() + Buffer->size();
45a16fe65bSRui Ueyama   }
46a16fe65bSRui Ueyama 
47a16fe65bSRui Ueyama   size_t getBufferSize() const override { return Buffer->size(); }
48a16fe65bSRui Ueyama 
490d7a38a8SRafael Espindola   Error commit() override {
50a16fe65bSRui Ueyama     // Unmap buffer, letting OS flush dirty pages to file on disk.
51a16fe65bSRui Ueyama     Buffer.reset();
52a16fe65bSRui Ueyama 
53a16fe65bSRui Ueyama     // Atomically replace the existing file with the new one.
54a16fe65bSRui Ueyama     auto EC = fs::rename(TempPath, FinalPath);
55a16fe65bSRui Ueyama     sys::DontRemoveFileOnSignal(TempPath);
560d7a38a8SRafael Espindola     return errorCodeToError(EC);
57a16fe65bSRui Ueyama   }
58a16fe65bSRui Ueyama 
59a16fe65bSRui Ueyama   ~OnDiskBuffer() override {
601a4398a1SReid Kleckner     // Close the mapping before deleting the temp file, so that the removal
611a4398a1SReid Kleckner     // succeeds.
62a16fe65bSRui Ueyama     Buffer.reset();
63a16fe65bSRui Ueyama     fs::remove(TempPath);
645fce8c4fSNick Kledzik   }
655fce8c4fSNick Kledzik 
66a16fe65bSRui Ueyama private:
67a16fe65bSRui Ueyama   std::unique_ptr<fs::mapped_file_region> Buffer;
68a16fe65bSRui Ueyama   std::string TempPath;
69a16fe65bSRui Ueyama };
70a16fe65bSRui Ueyama 
71a16fe65bSRui Ueyama // A FileOutputBuffer which keeps data in memory and writes to the final
72a16fe65bSRui Ueyama // output file on commit(). This is used only when we cannot use OnDiskBuffer.
73a16fe65bSRui Ueyama class InMemoryBuffer : public FileOutputBuffer {
74a16fe65bSRui Ueyama public:
75a16fe65bSRui Ueyama   InMemoryBuffer(StringRef Path, MemoryBlock Buf, unsigned Mode)
76a16fe65bSRui Ueyama       : FileOutputBuffer(Path), Buffer(Buf), Mode(Mode) {}
77a16fe65bSRui Ueyama 
78a16fe65bSRui Ueyama   uint8_t *getBufferStart() const override { return (uint8_t *)Buffer.base(); }
79a16fe65bSRui Ueyama 
80a16fe65bSRui Ueyama   uint8_t *getBufferEnd() const override {
81a16fe65bSRui Ueyama     return (uint8_t *)Buffer.base() + Buffer.size();
82a16fe65bSRui Ueyama   }
83a16fe65bSRui Ueyama 
84a16fe65bSRui Ueyama   size_t getBufferSize() const override { return Buffer.size(); }
85a16fe65bSRui Ueyama 
860d7a38a8SRafael Espindola   Error commit() override {
87d4b24edaSRafael Espindola     int FD;
88a16fe65bSRui Ueyama     std::error_code EC;
89a16fe65bSRui Ueyama     if (auto EC = openFileForWrite(FinalPath, FD, fs::F_None, Mode))
900d7a38a8SRafael Espindola       return errorCodeToError(EC);
91a16fe65bSRui Ueyama     raw_fd_ostream OS(FD, /*shouldClose=*/true, /*unbuffered=*/true);
92a16fe65bSRui Ueyama     OS << StringRef((const char *)Buffer.base(), Buffer.size());
930d7a38a8SRafael Espindola     return Error::success();
94d4b24edaSRafael Espindola   }
95d4b24edaSRafael Espindola 
96a16fe65bSRui Ueyama private:
97a16fe65bSRui Ueyama   OwningMemoryBlock Buffer;
98a16fe65bSRui Ueyama   unsigned Mode;
99a16fe65bSRui Ueyama };
100a16fe65bSRui Ueyama 
101*f6490e04SRui Ueyama static Expected<std::unique_ptr<InMemoryBuffer>>
102*f6490e04SRui Ueyama createInMemoryBuffer(StringRef Path, size_t Size, unsigned Mode) {
103*f6490e04SRui Ueyama   std::error_code EC;
104*f6490e04SRui Ueyama   MemoryBlock MB = Memory::allocateMappedMemory(
105*f6490e04SRui Ueyama       Size, nullptr, sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC);
106*f6490e04SRui Ueyama   if (EC)
107*f6490e04SRui Ueyama     return errorCodeToError(EC);
108*f6490e04SRui Ueyama   return llvm::make_unique<InMemoryBuffer>(Path, MB, Mode);
109*f6490e04SRui Ueyama }
110*f6490e04SRui Ueyama 
111*f6490e04SRui Ueyama static Expected<std::unique_ptr<OnDiskBuffer>>
112*f6490e04SRui Ueyama createOnDiskBuffer(StringRef Path, size_t Size, unsigned Mode) {
113a16fe65bSRui Ueyama   // Create new file in same directory but with random name.
114a16fe65bSRui Ueyama   SmallString<128> TempPath;
115a16fe65bSRui Ueyama   int FD;
116a16fe65bSRui Ueyama   if (auto EC = fs::createUniqueFile(Path + ".tmp%%%%%%%", FD, TempPath, Mode))
117e0df357dSRafael Espindola     return errorCodeToError(EC);
1185fce8c4fSNick Kledzik 
119a16fe65bSRui Ueyama   sys::RemoveFileOnSignal(TempPath);
1207dbb5778SRafael Espindola 
121da9bc2e5SRui Ueyama #ifndef LLVM_ON_WIN32
122da9bc2e5SRui Ueyama   // On Windows, CreateFileMapping (the mmap function on Windows)
123da9bc2e5SRui Ueyama   // automatically extends the underlying file. We don't need to
124da9bc2e5SRui Ueyama   // extend the file beforehand. _chsize (ftruncate on Windows) is
125da9bc2e5SRui Ueyama   // pretty slow just like it writes specified amount of bytes,
126a16fe65bSRui Ueyama   // so we should avoid calling that function.
127a16fe65bSRui Ueyama   if (auto EC = fs::resize_file(FD, Size))
128e0df357dSRafael Espindola     return errorCodeToError(EC);
129da9bc2e5SRui Ueyama #endif
130c69f13bfSRafael Espindola 
131a16fe65bSRui Ueyama   // Mmap it.
132a16fe65bSRui Ueyama   std::error_code EC;
133a16fe65bSRui Ueyama   auto MappedFile = llvm::make_unique<fs::mapped_file_region>(
134a16fe65bSRui Ueyama       FD, fs::mapped_file_region::readwrite, Size, 0, EC);
135a16fe65bSRui Ueyama   close(FD);
1365fce8c4fSNick Kledzik   if (EC)
137e0df357dSRafael Espindola     return errorCodeToError(EC);
138a16fe65bSRui Ueyama   return llvm::make_unique<OnDiskBuffer>(Path, TempPath, std::move(MappedFile));
1395fce8c4fSNick Kledzik }
1405fce8c4fSNick Kledzik 
141a16fe65bSRui Ueyama // Create an instance of FileOutputBuffer.
142e0df357dSRafael Espindola Expected<std::unique_ptr<FileOutputBuffer>>
143a16fe65bSRui Ueyama FileOutputBuffer::create(StringRef Path, size_t Size, unsigned Flags) {
144a16fe65bSRui Ueyama   unsigned Mode = fs::all_read | fs::all_write;
145a16fe65bSRui Ueyama   if (Flags & F_executable)
146a16fe65bSRui Ueyama     Mode |= fs::all_exe;
1475fce8c4fSNick Kledzik 
148a16fe65bSRui Ueyama   fs::file_status Stat;
149a16fe65bSRui Ueyama   fs::status(Path, Stat);
150d4b24edaSRafael Espindola 
151a16fe65bSRui Ueyama   // Usually, we want to create OnDiskBuffer to create a temporary file in
152a16fe65bSRui Ueyama   // the same directory as the destination file and atomically replaces it
153a16fe65bSRui Ueyama   // by rename(2).
154a16fe65bSRui Ueyama   //
155a16fe65bSRui Ueyama   // However, if the destination file is a special file, we don't want to
156a16fe65bSRui Ueyama   // use rename (e.g. we don't want to replace /dev/null with a regular
157a16fe65bSRui Ueyama   // file.) If that's the case, we create an in-memory buffer, open the
158a16fe65bSRui Ueyama   // destination file and write to it on commit().
159a16fe65bSRui Ueyama   switch (Stat.type()) {
160a16fe65bSRui Ueyama   case fs::file_type::directory_file:
161e0df357dSRafael Espindola     return errorCodeToError(errc::is_a_directory);
162a16fe65bSRui Ueyama   case fs::file_type::regular_file:
163a16fe65bSRui Ueyama   case fs::file_type::file_not_found:
164a16fe65bSRui Ueyama   case fs::file_type::status_error:
165*f6490e04SRui Ueyama     return createOnDiskBuffer(Path, Size, Mode);
166a16fe65bSRui Ueyama   default:
167*f6490e04SRui Ueyama     return createInMemoryBuffer(Path, Size, Mode);
1685fce8c4fSNick Kledzik   }
169a16fe65bSRui Ueyama }
170