Invoke a worker thread to pre-populate file contents

main.cc

@@ -751,9 +751,6 @@ int main(int argc, char **argv) {
   std::vector<std::string_view> file_args;
   config = parse_nonpositional_args(arg_vector, file_args);
 
-  if (config.output == "")
-    Fatal() << "-o option is missing";
-
   if (!config.preload)
     if (i64 code; resume_daemon(argv, &code))
       exit(code);
@@ -777,6 +774,10 @@ int main(int argc, char **argv) {
     on_complete = fork_child();
   }
 
+  if (config.output == "")
+    Fatal() << "-o option is missing";
+  OutputFile *output_file = OutputFile::open(config.output);
+
   if (config.pic)
     config.image_base = 0;
 
@@ -1062,8 +1063,10 @@ int main(int argc, char **argv) {
   t_before_copy.stop();
 
   // Create an output file
-  OutputFile *file = OutputFile::open(config.output, filesize);
-  out::buf = file->buf;
+  out::buf = output_file->get_buffer(filesize);
+
+  if (config.filler != -1)
+    memset(out::buf, config.filler, filesize);
 
   Timer t_copy("copy");
 
@@ -1085,7 +1088,7 @@ int main(int argc, char **argv) {
     out::buildid->write_buildid(filesize);
   }
 
-  file->close();
+  std::function<void()> closer = output_file->close();
 
   t_copy.stop();
   t_total.stop();
@@ -1106,6 +1109,7 @@ int main(int argc, char **argv) {
   if (on_complete)
     on_complete();
 
+  closer();
   if (config.quick_exit)
     std::quick_exit(0);
   return 0;

mold.h

@@ -1056,17 +1056,21 @@ void parse_version_script(std::string path);
 
 class OutputFile {
 public:
-  static OutputFile *open(std::string path, u64 filesize);
-  virtual void close() = 0;
+  static OutputFile *open(std::string path);
+
+  virtual u8 *get_buffer(u64 filesize) = 0;
+  virtual std::function<void()> close() = 0;
 
-  u8 *buf;
   static inline char *tmpfile;
 
 protected:
-  OutputFile(std::string path, u64 filesize) : path(path), filesize(filesize) {}
+  OutputFile(std::string path) : path(path) {}
 
   std::string path;
-  u64 filesize;
+  int fd = -1;
+  u64 filesize = -1;
+  u64 populate_size = -1;
+  u8 *buf = nullptr;
 };
 
 //

output_file.cc

@@ -5,6 +5,7 @@
 #include <sys/mman.h>
 #include <sys/stat.h>
 #include <sys/types.h>
+#include <tbb/task_group.h>
 
 static u32 get_umask() {
   u32 orig_umask = umask(0);
@@ -14,59 +15,101 @@ static u32 get_umask() {
 
 class MemoryMappedOutputFile : public OutputFile {
 public:
-  MemoryMappedOutputFile(std::string path, i64 filesize)
-    : OutputFile(path, filesize) {
+  MemoryMappedOutputFile(std::string path) : OutputFile(path) {
     std::string dir = dirname(strdup(config.output.c_str()));
     tmpfile = strdup((dir + "/.mold-XXXXXX").c_str());
-    i64 fd = mkstemp(tmpfile);
+    fd = mkstemp(tmpfile);
     if (fd == -1)
       Fatal() << "cannot open " << tmpfile <<  ": " << strerror(errno);
-
-    if (rename(config.output.c_str(), tmpfile) == 0) {
-      ::close(fd);
-      fd = ::open(tmpfile, O_RDWR | O_CREAT, 0777);
-      if (fd == -1) {
-        if (errno != ETXTBSY)
-          Fatal() << "cannot open " << config.output << ": " << strerror(errno);
-        unlink(tmpfile);
-        fd = ::open(tmpfile, O_RDWR | O_CREAT, 0777);
-        if (fd == -1)
-          Fatal() << "cannot open " << config.output << ": " << strerror(errno);
-      }
-    }
-
-    if (ftruncate(fd, filesize))
-      Fatal() << "ftruncate failed";
-
     if (fchmod(fd, (0777 & ~get_umask())) == -1)
-      Fatal() << "fchmod failed";
-
-    buf = (u8 *)mmap(nullptr, filesize, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
-    if (buf == MAP_FAILED)
-      Fatal() << config.output << ": mmap failed: " << strerror(errno);
-    ::close(fd);
+      Fatal() << tmpfile << ": fchmod failed";
   }
 
-  void close() override {
-    Timer t("munmap");
+  void populate(u64 size) {
+    populate_size = size;
+
+    if (ftruncate(fd, size))
+      Fatal() << "ftruncate failed";
+
+    buf = (u8 *)mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+    if (buf == MAP_FAILED)
+      Fatal() << config.output << ": mmap failed: " << strerror(errno);
+
+    tg.run([=]() { memset(buf, 0, size); });
+  }
+
+  u8 *get_buffer(u64 size) override {
+    Timer t("get_buffer");
+    tg.wait();
+
+    filesize = size;
+    if (ftruncate(fd, size))
+      Fatal() << "ftruncate failed";
+
+    if (buf) {
+      buf = (u8 *)mremap(buf, populate_size, size, MREMAP_MAYMOVE);
+      if (buf == MAP_FAILED)
+        Fatal() << config.output << ": mremap failed: " << strerror(errno);
+    } else {
+      buf = (u8 *)mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+      if (buf == MAP_FAILED)
+        Fatal() << config.output << ": mmap failed: " << strerror(errno);
+    }
+    ::close(fd);
+    return buf;
+  }
+
+  std::function<void()> close() override {
+    Timer t("close");
     munmap(buf, filesize);
+
+    char *existing = rename_existing_file();
     if (rename(tmpfile, config.output.c_str()) == -1)
       Fatal() << config.output << ": rename filed: " << strerror(errno);
     tmpfile = nullptr;
+
+    if (existing)
+      return [=]() { unlink(existing); };
+    return []() {};
   }
+
+private:
+  char *rename_existing_file() {
+    Timer t("rename_existing_file");
+    char *path = strdup(config.output.c_str());
+    struct stat st;
+    if (stat(path, &st) != 0)
+      return nullptr;
+
+    std::string dir = dirname(strdup(path));
+    char *tmp = strdup((dir + "/.mold-XXXXXX").c_str());
+    int fd = mkstemp(tmp);
+    if (fd == -1)
+      return nullptr;
+    ::close(fd);
+
+    if (::rename(path, tmp) == -1)
+      return nullptr;
+    return tmp;
+  }
+
+  tbb::task_group tg;
 };
 
 class MallocOutputFile : public OutputFile {
 public:
-  MallocOutputFile(std::string path, u64 filesize)
-    : OutputFile(path, filesize) {
-    buf = (u8 *)mmap(NULL, filesize, PROT_READ | PROT_WRITE,
+  MallocOutputFile(std::string path) : OutputFile(path) {}
+
+  u8 *get_buffer(u64 size) override {
+    filesize = size;
+    buf = (u8 *)mmap(NULL, size, PROT_READ | PROT_WRITE,
                      MAP_SHARED | MAP_ANONYMOUS, -1, 0);
     if (buf == MAP_FAILED)
       Fatal() << "mmap failed: " << strerror(errno);
+    return buf;
   }
 
-  void close() override {
+  std::function<void()> close() override {
     Timer t("munmap");
     i64 fd = ::open(path.c_str(), O_RDWR | O_CREAT, 0777);
     if (fd == -1)
@@ -75,24 +118,22 @@ public:
     FILE *fp = fdopen(fd, "w");
     fwrite(buf, filesize, 1, fp);
     fclose(fp);
+    return []() {};
   }
 };
 
-OutputFile *OutputFile::open(std::string path, u64 filesize) {
+OutputFile *OutputFile::open(std::string path) {
   Timer t("open_file");
+  i64 size = -1;
 
-  bool is_special = false;
-  struct stat st;
-  if (stat(path.c_str(), &st) == 0 && (st.st_mode & S_IFMT) != S_IFREG)
-    is_special = true;
+  if (struct stat st; stat(path.c_str(), &st) == 0) {
+    if ((st.st_mode & S_IFMT) != S_IFREG)
+      return new MallocOutputFile(path);
+    size = st.st_size;
+  }
 
-  OutputFile *file;
-  if (is_special)
-    file = new MallocOutputFile(path, filesize);
-  else
-    file = new MemoryMappedOutputFile(path, filesize);
-
-  if (config.filler != -1)
-    memset(file->buf, config.filler, filesize);
+  auto *file = new MemoryMappedOutputFile(path);
+  if (size != -1)
+    file->populate(size);
   return file;
 }