15fce8c4fSNick Kledzik //===- FileOutputBuffer.cpp - File Output Buffer ----------------*- C++ -*-===//
25fce8c4fSNick Kledzik //
3*2946cd70SChandler Carruth // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4*2946cd70SChandler Carruth // See https://llvm.org/LICENSE.txt for license information.
5*2946cd70SChandler Carruth // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
65fce8c4fSNick Kledzik //
75fce8c4fSNick Kledzik //===----------------------------------------------------------------------===//
85fce8c4fSNick Kledzik //
95fce8c4fSNick Kledzik // Utility for creating a in-memory buffer that will be written to a file.
105fce8c4fSNick Kledzik //
115fce8c4fSNick Kledzik //===----------------------------------------------------------------------===//
125fce8c4fSNick Kledzik 
13d9903888SChandler Carruth #include "llvm/Support/FileOutputBuffer.h"
1416132e6fSBenjamin Kramer #include "llvm/ADT/STLExtras.h"
1516132e6fSBenjamin Kramer #include "llvm/ADT/SmallString.h"
1616132e6fSBenjamin Kramer #include "llvm/Support/Errc.h"
17a16fe65bSRui Ueyama #include "llvm/Support/Memory.h"
18d4b24edaSRafael Espindola #include "llvm/Support/Path.h"
19a6e9c3e4SRafael Espindola #include <system_error>
205fce8c4fSNick Kledzik 
217eb1f185SRafael Espindola #if !defined(_MSC_VER) && !defined(__MINGW32__)
227eb1f185SRafael Espindola #include <unistd.h>
237eb1f185SRafael Espindola #else
247eb1f185SRafael Espindola #include <io.h>
257eb1f185SRafael Espindola #endif
267eb1f185SRafael Espindola 
27a16fe65bSRui Ueyama using namespace llvm;
28a16fe65bSRui Ueyama using namespace llvm::sys;
295fce8c4fSNick Kledzik 
3051ebcaafSBenjamin Kramer namespace {
31a16fe65bSRui Ueyama // A FileOutputBuffer which creates a temporary file in the same directory
32a16fe65bSRui Ueyama // as the final output file. The final output file is atomically replaced
33a16fe65bSRui Ueyama // with the temporary file on commit().
34a16fe65bSRui Ueyama class OnDiskBuffer : public FileOutputBuffer {
35a16fe65bSRui Ueyama public:
3658fe67a9SRafael Espindola   OnDiskBuffer(StringRef Path, fs::TempFile Temp,
37a16fe65bSRui Ueyama                std::unique_ptr<fs::mapped_file_region> Buf)
3858fe67a9SRafael Espindola       : FileOutputBuffer(Path), Buffer(std::move(Buf)), Temp(std::move(Temp)) {}
395fce8c4fSNick Kledzik 
40a16fe65bSRui Ueyama   uint8_t *getBufferStart() const override { return (uint8_t *)Buffer->data(); }
41a16fe65bSRui Ueyama 
42a16fe65bSRui Ueyama   uint8_t *getBufferEnd() const override {
43a16fe65bSRui Ueyama     return (uint8_t *)Buffer->data() + Buffer->size();
44a16fe65bSRui Ueyama   }
45a16fe65bSRui Ueyama 
46a16fe65bSRui Ueyama   size_t getBufferSize() const override { return Buffer->size(); }
47a16fe65bSRui Ueyama 
480d7a38a8SRafael Espindola   Error commit() override {
49a16fe65bSRui Ueyama     // Unmap buffer, letting OS flush dirty pages to file on disk.
50a16fe65bSRui Ueyama     Buffer.reset();
51a16fe65bSRui Ueyama 
52a16fe65bSRui Ueyama     // Atomically replace the existing file with the new one.
5358fe67a9SRafael Espindola     return Temp.keep(FinalPath);
54a16fe65bSRui Ueyama   }
55a16fe65bSRui Ueyama 
56a16fe65bSRui Ueyama   ~OnDiskBuffer() override {
571a4398a1SReid Kleckner     // Close the mapping before deleting the temp file, so that the removal
581a4398a1SReid Kleckner     // succeeds.
59a16fe65bSRui Ueyama     Buffer.reset();
6058fe67a9SRafael Espindola     consumeError(Temp.discard());
615fce8c4fSNick Kledzik   }
625fce8c4fSNick Kledzik 
634153e9fbSMartin Storsjo   void discard() override {
644153e9fbSMartin Storsjo     // Delete the temp file if it still was open, but keeping the mapping
654153e9fbSMartin Storsjo     // active.
664153e9fbSMartin Storsjo     consumeError(Temp.discard());
674153e9fbSMartin Storsjo   }
684153e9fbSMartin Storsjo 
69a16fe65bSRui Ueyama private:
70a16fe65bSRui Ueyama   std::unique_ptr<fs::mapped_file_region> Buffer;
7158fe67a9SRafael Espindola   fs::TempFile Temp;
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 
81a16fe65bSRui Ueyama   uint8_t *getBufferStart() const override { return (uint8_t *)Buffer.base(); }
82a16fe65bSRui Ueyama 
83a16fe65bSRui Ueyama   uint8_t *getBufferEnd() const override {
84a16fe65bSRui Ueyama     return (uint8_t *)Buffer.base() + Buffer.size();
85a16fe65bSRui Ueyama   }
86a16fe65bSRui Ueyama 
87a16fe65bSRui Ueyama   size_t getBufferSize() const override { return Buffer.size(); }
88a16fe65bSRui Ueyama 
890d7a38a8SRafael Espindola   Error commit() override {
901f67a3cbSZachary Turner     using namespace sys::fs;
91d4b24edaSRafael Espindola     int FD;
92a16fe65bSRui Ueyama     std::error_code EC;
93154a72ddSZachary Turner     if (auto EC =
94154a72ddSZachary Turner             openFileForWrite(FinalPath, FD, CD_CreateAlways, OF_None, Mode))
950d7a38a8SRafael Espindola       return errorCodeToError(EC);
96a16fe65bSRui Ueyama     raw_fd_ostream OS(FD, /*shouldClose=*/true, /*unbuffered=*/true);
97a16fe65bSRui Ueyama     OS << StringRef((const char *)Buffer.base(), Buffer.size());
980d7a38a8SRafael Espindola     return Error::success();
99d4b24edaSRafael Espindola   }
100d4b24edaSRafael Espindola 
101a16fe65bSRui Ueyama private:
102a16fe65bSRui Ueyama   OwningMemoryBlock Buffer;
103a16fe65bSRui Ueyama   unsigned Mode;
104a16fe65bSRui Ueyama };
10551ebcaafSBenjamin Kramer } // namespace
106a16fe65bSRui Ueyama 
107f6490e04SRui Ueyama static Expected<std::unique_ptr<InMemoryBuffer>>
108f6490e04SRui Ueyama createInMemoryBuffer(StringRef Path, size_t Size, unsigned Mode) {
109f6490e04SRui Ueyama   std::error_code EC;
110f6490e04SRui Ueyama   MemoryBlock MB = Memory::allocateMappedMemory(
111f6490e04SRui Ueyama       Size, nullptr, sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC);
112f6490e04SRui Ueyama   if (EC)
113f6490e04SRui Ueyama     return errorCodeToError(EC);
114f6490e04SRui Ueyama   return llvm::make_unique<InMemoryBuffer>(Path, MB, Mode);
115f6490e04SRui Ueyama }
116f6490e04SRui Ueyama 
117f6490e04SRui Ueyama static Expected<std::unique_ptr<OnDiskBuffer>>
1188e7600dcSRui Ueyama createOnDiskBuffer(StringRef Path, size_t Size, unsigned Mode) {
11958fe67a9SRafael Espindola   Expected<fs::TempFile> FileOrErr =
12058fe67a9SRafael Espindola       fs::TempFile::create(Path + ".tmp%%%%%%%", Mode);
12158fe67a9SRafael Espindola   if (!FileOrErr)
12258fe67a9SRafael Espindola     return FileOrErr.takeError();
12358fe67a9SRafael Espindola   fs::TempFile File = std::move(*FileOrErr);
1247dbb5778SRafael Espindola 
125712e8d29SNico Weber #ifndef _WIN32
126da9bc2e5SRui Ueyama   // On Windows, CreateFileMapping (the mmap function on Windows)
127da9bc2e5SRui Ueyama   // automatically extends the underlying file. We don't need to
128da9bc2e5SRui Ueyama   // extend the file beforehand. _chsize (ftruncate on Windows) is
129da9bc2e5SRui Ueyama   // pretty slow just like it writes specified amount of bytes,
130a16fe65bSRui Ueyama   // so we should avoid calling that function.
13158fe67a9SRafael Espindola   if (auto EC = fs::resize_file(File.FD, Size)) {
13258fe67a9SRafael Espindola     consumeError(File.discard());
133e0df357dSRafael Espindola     return errorCodeToError(EC);
13458fe67a9SRafael Espindola   }
135da9bc2e5SRui Ueyama #endif
136c69f13bfSRafael Espindola 
137a16fe65bSRui Ueyama   // Mmap it.
138a16fe65bSRui Ueyama   std::error_code EC;
139a16fe65bSRui Ueyama   auto MappedFile = llvm::make_unique<fs::mapped_file_region>(
14058fe67a9SRafael Espindola       File.FD, fs::mapped_file_region::readwrite, Size, 0, EC);
14158fe67a9SRafael Espindola   if (EC) {
14258fe67a9SRafael Espindola     consumeError(File.discard());
143e0df357dSRafael Espindola     return errorCodeToError(EC);
14458fe67a9SRafael Espindola   }
14558fe67a9SRafael Espindola   return llvm::make_unique<OnDiskBuffer>(Path, std::move(File),
14658fe67a9SRafael Espindola                                          std::move(MappedFile));
1475fce8c4fSNick Kledzik }
1485fce8c4fSNick Kledzik 
149a16fe65bSRui Ueyama // Create an instance of FileOutputBuffer.
150e0df357dSRafael Espindola Expected<std::unique_ptr<FileOutputBuffer>>
151a16fe65bSRui Ueyama FileOutputBuffer::create(StringRef Path, size_t Size, unsigned Flags) {
152a16fe65bSRui Ueyama   unsigned Mode = fs::all_read | fs::all_write;
153a16fe65bSRui Ueyama   if (Flags & F_executable)
154a16fe65bSRui Ueyama     Mode |= fs::all_exe;
1555fce8c4fSNick Kledzik 
156a16fe65bSRui Ueyama   fs::file_status Stat;
157a16fe65bSRui Ueyama   fs::status(Path, Stat);
158d4b24edaSRafael Espindola 
159a16fe65bSRui Ueyama   // Usually, we want to create OnDiskBuffer to create a temporary file in
160a16fe65bSRui Ueyama   // the same directory as the destination file and atomically replaces it
161a16fe65bSRui Ueyama   // by rename(2).
162a16fe65bSRui Ueyama   //
163a16fe65bSRui Ueyama   // However, if the destination file is a special file, we don't want to
164a16fe65bSRui Ueyama   // use rename (e.g. we don't want to replace /dev/null with a regular
165a16fe65bSRui Ueyama   // file.) If that's the case, we create an in-memory buffer, open the
166a16fe65bSRui Ueyama   // destination file and write to it on commit().
167a16fe65bSRui Ueyama   switch (Stat.type()) {
168a16fe65bSRui Ueyama   case fs::file_type::directory_file:
169e0df357dSRafael Espindola     return errorCodeToError(errc::is_a_directory);
170a16fe65bSRui Ueyama   case fs::file_type::regular_file:
171a16fe65bSRui Ueyama   case fs::file_type::file_not_found:
172a16fe65bSRui Ueyama   case fs::file_type::status_error:
1738e7600dcSRui Ueyama     return createOnDiskBuffer(Path, Size, Mode);
174a16fe65bSRui Ueyama   default:
175f6490e04SRui Ueyama     return createInMemoryBuffer(Path, Size, Mode);
1765fce8c4fSNick Kledzik   }
177a16fe65bSRui Ueyama }
178