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"
20a6e9c3e4SRafael Espindola #include <system_error>
215fce8c4fSNick Kledzik 
227eb1f185SRafael Espindola #if !defined(_MSC_VER) && !defined(__MINGW32__)
237eb1f185SRafael Espindola #include <unistd.h>
247eb1f185SRafael Espindola #else
257eb1f185SRafael Espindola #include <io.h>
267eb1f185SRafael Espindola #endif
277eb1f185SRafael Espindola 
28a16fe65bSRui Ueyama using namespace llvm;
29a16fe65bSRui Ueyama using namespace llvm::sys;
305fce8c4fSNick Kledzik 
3151ebcaafSBenjamin Kramer namespace {
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:
3758fe67a9SRafael Espindola   OnDiskBuffer(StringRef Path, fs::TempFile Temp,
38a16fe65bSRui Ueyama                std::unique_ptr<fs::mapped_file_region> Buf)
3958fe67a9SRafael Espindola       : FileOutputBuffer(Path), Buffer(std::move(Buf)), Temp(std::move(Temp)) {}
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.
5458fe67a9SRafael Espindola     return Temp.keep(FinalPath);
55a16fe65bSRui Ueyama   }
56a16fe65bSRui Ueyama 
57a16fe65bSRui Ueyama   ~OnDiskBuffer() override {
581a4398a1SReid Kleckner     // Close the mapping before deleting the temp file, so that the removal
591a4398a1SReid Kleckner     // succeeds.
60a16fe65bSRui Ueyama     Buffer.reset();
6158fe67a9SRafael Espindola     consumeError(Temp.discard());
625fce8c4fSNick Kledzik   }
635fce8c4fSNick Kledzik 
644153e9fbSMartin Storsjo   void discard() override {
654153e9fbSMartin Storsjo     // Delete the temp file if it still was open, but keeping the mapping
664153e9fbSMartin Storsjo     // active.
674153e9fbSMartin Storsjo     consumeError(Temp.discard());
684153e9fbSMartin Storsjo   }
694153e9fbSMartin Storsjo 
70a16fe65bSRui Ueyama private:
71a16fe65bSRui Ueyama   std::unique_ptr<fs::mapped_file_region> Buffer;
7258fe67a9SRafael Espindola   fs::TempFile Temp;
73a16fe65bSRui Ueyama };
74a16fe65bSRui Ueyama 
75a16fe65bSRui Ueyama // A FileOutputBuffer which keeps data in memory and writes to the final
76a16fe65bSRui Ueyama // output file on commit(). This is used only when we cannot use OnDiskBuffer.
77a16fe65bSRui Ueyama class InMemoryBuffer : public FileOutputBuffer {
78a16fe65bSRui Ueyama public:
79a16fe65bSRui Ueyama   InMemoryBuffer(StringRef Path, MemoryBlock Buf, unsigned Mode)
80a16fe65bSRui Ueyama       : FileOutputBuffer(Path), Buffer(Buf), Mode(Mode) {}
81a16fe65bSRui Ueyama 
82a16fe65bSRui Ueyama   uint8_t *getBufferStart() const override { return (uint8_t *)Buffer.base(); }
83a16fe65bSRui Ueyama 
84a16fe65bSRui Ueyama   uint8_t *getBufferEnd() const override {
85a16fe65bSRui Ueyama     return (uint8_t *)Buffer.base() + Buffer.size();
86a16fe65bSRui Ueyama   }
87a16fe65bSRui Ueyama 
88a16fe65bSRui Ueyama   size_t getBufferSize() const override { return Buffer.size(); }
89a16fe65bSRui Ueyama 
900d7a38a8SRafael Espindola   Error commit() override {
911f67a3cbSZachary Turner     using namespace sys::fs;
92d4b24edaSRafael Espindola     int FD;
93a16fe65bSRui Ueyama     std::error_code EC;
94154a72ddSZachary Turner     if (auto EC =
95154a72ddSZachary Turner             openFileForWrite(FinalPath, FD, CD_CreateAlways, OF_None, Mode))
960d7a38a8SRafael Espindola       return errorCodeToError(EC);
97a16fe65bSRui Ueyama     raw_fd_ostream OS(FD, /*shouldClose=*/true, /*unbuffered=*/true);
98a16fe65bSRui Ueyama     OS << StringRef((const char *)Buffer.base(), Buffer.size());
990d7a38a8SRafael Espindola     return Error::success();
100d4b24edaSRafael Espindola   }
101d4b24edaSRafael Espindola 
102a16fe65bSRui Ueyama private:
103a16fe65bSRui Ueyama   OwningMemoryBlock Buffer;
104a16fe65bSRui Ueyama   unsigned Mode;
105a16fe65bSRui Ueyama };
10651ebcaafSBenjamin Kramer } // namespace
107a16fe65bSRui Ueyama 
108f6490e04SRui Ueyama static Expected<std::unique_ptr<InMemoryBuffer>>
109f6490e04SRui Ueyama createInMemoryBuffer(StringRef Path, size_t Size, unsigned Mode) {
110f6490e04SRui Ueyama   std::error_code EC;
111f6490e04SRui Ueyama   MemoryBlock MB = Memory::allocateMappedMemory(
112f6490e04SRui Ueyama       Size, nullptr, sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC);
113f6490e04SRui Ueyama   if (EC)
114f6490e04SRui Ueyama     return errorCodeToError(EC);
115f6490e04SRui Ueyama   return llvm::make_unique<InMemoryBuffer>(Path, MB, Mode);
116f6490e04SRui Ueyama }
117f6490e04SRui Ueyama 
118f6490e04SRui Ueyama static Expected<std::unique_ptr<OnDiskBuffer>>
119*8e7600dcSRui Ueyama createOnDiskBuffer(StringRef Path, size_t Size, unsigned Mode) {
12058fe67a9SRafael Espindola   Expected<fs::TempFile> FileOrErr =
12158fe67a9SRafael Espindola       fs::TempFile::create(Path + ".tmp%%%%%%%", Mode);
12258fe67a9SRafael Espindola   if (!FileOrErr)
12358fe67a9SRafael Espindola     return FileOrErr.takeError();
12458fe67a9SRafael Espindola   fs::TempFile File = std::move(*FileOrErr);
1257dbb5778SRafael Espindola 
126712e8d29SNico Weber #ifndef _WIN32
127da9bc2e5SRui Ueyama   // On Windows, CreateFileMapping (the mmap function on Windows)
128da9bc2e5SRui Ueyama   // automatically extends the underlying file. We don't need to
129da9bc2e5SRui Ueyama   // extend the file beforehand. _chsize (ftruncate on Windows) is
130da9bc2e5SRui Ueyama   // pretty slow just like it writes specified amount of bytes,
131a16fe65bSRui Ueyama   // so we should avoid calling that function.
13258fe67a9SRafael Espindola   if (auto EC = fs::resize_file(File.FD, Size)) {
13358fe67a9SRafael Espindola     consumeError(File.discard());
134e0df357dSRafael Espindola     return errorCodeToError(EC);
13558fe67a9SRafael Espindola   }
136da9bc2e5SRui Ueyama #endif
137c69f13bfSRafael Espindola 
138a16fe65bSRui Ueyama   // Mmap it.
139a16fe65bSRui Ueyama   std::error_code EC;
140a16fe65bSRui Ueyama   auto MappedFile = llvm::make_unique<fs::mapped_file_region>(
14158fe67a9SRafael Espindola       File.FD, fs::mapped_file_region::readwrite, Size, 0, EC);
14258fe67a9SRafael Espindola   if (EC) {
14358fe67a9SRafael Espindola     consumeError(File.discard());
144e0df357dSRafael Espindola     return errorCodeToError(EC);
14558fe67a9SRafael Espindola   }
14658fe67a9SRafael Espindola   return llvm::make_unique<OnDiskBuffer>(Path, std::move(File),
14758fe67a9SRafael Espindola                                          std::move(MappedFile));
1485fce8c4fSNick Kledzik }
1495fce8c4fSNick Kledzik 
150a16fe65bSRui Ueyama // Create an instance of FileOutputBuffer.
151e0df357dSRafael Espindola Expected<std::unique_ptr<FileOutputBuffer>>
152a16fe65bSRui Ueyama FileOutputBuffer::create(StringRef Path, size_t Size, unsigned Flags) {
153a16fe65bSRui Ueyama   unsigned Mode = fs::all_read | fs::all_write;
154a16fe65bSRui Ueyama   if (Flags & F_executable)
155a16fe65bSRui Ueyama     Mode |= fs::all_exe;
1565fce8c4fSNick Kledzik 
157a16fe65bSRui Ueyama   fs::file_status Stat;
158a16fe65bSRui Ueyama   fs::status(Path, Stat);
159d4b24edaSRafael Espindola 
160a16fe65bSRui Ueyama   // Usually, we want to create OnDiskBuffer to create a temporary file in
161a16fe65bSRui Ueyama   // the same directory as the destination file and atomically replaces it
162a16fe65bSRui Ueyama   // by rename(2).
163a16fe65bSRui Ueyama   //
164a16fe65bSRui Ueyama   // However, if the destination file is a special file, we don't want to
165a16fe65bSRui Ueyama   // use rename (e.g. we don't want to replace /dev/null with a regular
166a16fe65bSRui Ueyama   // file.) If that's the case, we create an in-memory buffer, open the
167a16fe65bSRui Ueyama   // destination file and write to it on commit().
168a16fe65bSRui Ueyama   switch (Stat.type()) {
169a16fe65bSRui Ueyama   case fs::file_type::directory_file:
170e0df357dSRafael Espindola     return errorCodeToError(errc::is_a_directory);
171a16fe65bSRui Ueyama   case fs::file_type::regular_file:
172a16fe65bSRui Ueyama   case fs::file_type::file_not_found:
173a16fe65bSRui Ueyama   case fs::file_type::status_error:
174*8e7600dcSRui Ueyama     return createOnDiskBuffer(Path, Size, Mode);
175a16fe65bSRui Ueyama   default:
176f6490e04SRui Ueyama     return createInMemoryBuffer(Path, Size, Mode);
1775fce8c4fSNick Kledzik   }
178a16fe65bSRui Ueyama }
179