[CAS/LazyMappedFileRegion] Fix racing issue during the initial creation of a database file

There's a racing issue when multiple processes try to create a new database file on the same directory.
Due to insufficient synchronization one process may try to use a newly created database file that another
process haven't finished initializing yet. This can manifest with errors like "database: bad magic" or
crashes.

rdar://104829275
(cherry picked from commit 4792257)

Author: akyrtzi

llvm/include/llvm/Support/FileSystem.h

@@ -1200,7 +1200,10 @@ std::error_code getRealPathFromHandle(file_t Handle,
 /// Lock the file.
 ///
 /// This function acts as @ref tryLockFile but it waits infinitely.
-std::error_code lockFile(int FD);
+/// \param FD file descriptor to use for locking.
+/// \param Exclusive if \p true use exclusive/writer lock, otherwise use
+/// shared/reader lock.
+std::error_code lockFile(int FD, bool Exclusive = true);
 
 /// Unlock the file.
 ///

llvm/lib/CAS/LazyMappedFileRegion.cpp

@@ -242,9 +242,59 @@ Expected<LazyMappedFileRegion> LazyMappedFileRegion::create(
 
   sys::fs::file_t File = sys::fs::convertFDToNativeFile(FD);
 
+  struct FileLockRAII {
+    LazyMappedFileRegion &LMFR;
+    bool IsLocked = false;
+
+    enum LockKind { Shared, Exclusive };
+
+    FileLockRAII(LazyMappedFileRegion &LMFR) : LMFR(LMFR) {}
+    ~FileLockRAII() { consumeError(unlock()); }
+
+    Error lock(LockKind LK) {
+      if (IsLocked)
+        return createStringError(inconvertibleErrorCode(),
+                                 LMFR.Path + " already locked");
+      if (std::error_code EC = sys::fs::lockFile(*LMFR.FD, LK == Exclusive))
+        return createFileError(LMFR.Path, EC);
+      IsLocked = true;
+      return Error::success();
+    }
+
+    Error unlock() {
+      if (IsLocked) {
+        IsLocked = false;
+        if (std::error_code EC = sys::fs::unlockFile(*LMFR.FD))
+          return createFileError(LMFR.Path, EC);
+      }
+      return Error::success();
+    }
+
+  } FileLock(LMFR);
+
+  // Use shared/reader locking in case another process is in the process of
+  // initializing the file.
+  if (Error E = FileLock.lock(FileLockRAII::Shared))
+    return std::move(E);
+
   sys::fs::file_status Status;
   if (std::error_code EC = sys::fs::status(File, Status))
     return errorCodeToError(EC);
+
+  if (Status.getSize() == 0) {
+    // Lock the file exclusively so only one process will do the initialization.
+    if (Error E = FileLock.unlock())
+      return std::move(E);
+    if (Error E = FileLock.lock(FileLockRAII::Exclusive))
+      return std::move(E);
+    if (std::error_code EC = sys::fs::status(File, Status))
+      return errorCodeToError(EC);
+  }
+
+  // At this point either the file is still empty (this process won the race to
+  // do the initialization) or we have the size for the completely initialized
+  // file.
+
   if (Status.getSize() > 0)
     // The file was already constructed.
     LMFR.CachedSize = Status.getSize();
@@ -263,11 +313,6 @@ Expected<LazyMappedFileRegion> LazyMappedFileRegion::create(
   if (!LMFR.IsConstructingNewFile)
     return std::move(LMFR);
 
-  // Lock the file so we can initialize it.
-  if (std::error_code EC = sys::fs::lockFile(*LMFR.FD))
-    return createFileError(Path, EC);
-  auto Unlock = make_scope_exit([FD = *LMFR.FD]() { sys::fs::unlockFile(FD); });
-
   // This is a new file. Resize to NewFileSize and run the constructor.
   if (Error E = NewFileConstructor(LMFR))
     return std::move(E);

llvm/lib/Support/Unix/Path.inc

@@ -1265,10 +1265,10 @@ std::error_code tryLockFile(int FD, std::chrono::milliseconds Timeout) {
   return make_error_code(errc::no_lock_available);
 }
 
-std::error_code lockFile(int FD) {
+std::error_code lockFile(int FD, bool Exclusive) {
   struct flock Lock;
   memset(&Lock, 0, sizeof(Lock));
-  Lock.l_type = F_WRLCK;
+  Lock.l_type = Exclusive ? F_WRLCK : F_RDLCK;
   Lock.l_whence = SEEK_SET;
   Lock.l_start = 0;
   Lock.l_len = 0;

llvm/lib/Support/Windows/Path.inc

@@ -1312,8 +1312,8 @@ std::error_code tryLockFile(int FD, std::chrono::milliseconds Timeout) {
   return mapWindowsError(ERROR_LOCK_VIOLATION);
 }
 
-std::error_code lockFile(int FD) {
-  DWORD Flags = LOCKFILE_EXCLUSIVE_LOCK;
+std::error_code lockFile(int FD, bool Exclusive) {
+  DWORD Flags = Exclusive ? LOCKFILE_EXCLUSIVE_LOCK : 0;
   OVERLAPPED OV = {};
   file_t File = convertFDToNativeFile(FD);
   if (::LockFileEx(File, Flags, 0, MAXDWORD, MAXDWORD, &OV))

llvm/test/CAS/Inputs/cas-creation-stress-test.sh

@@ -0,0 +1,23 @@
+#!/bin/sh
+
+# parameters are: <llvm-cas path> <CAS path> <file to ingest> <number of repeats>
+LLVM_CAS_TOOL=$1
+CAS_PATH=$2
+INGEST_FILE=$3
+NUM_REPEAT=$4
+
+set -e
+
+for c in $(seq 1 $NUM_REPEAT); do
+  rm -rf $CAS_PATH
+
+  pids=""
+  for x in $(seq 1 10); do
+    $LLVM_CAS_TOOL --ingest --data $INGEST_FILE --cas $CAS_PATH &
+    pids="$pids $!"
+  done
+
+  for pid in $pids; do
+    wait "$pid"
+  done
+done

llvm/test/CAS/databasefile-concurrent-creation.c

@@ -0,0 +1,10 @@
+// REQUIRES: ondisk_cas, shell
+
+// RUN: mkdir -p %t
+
+// This uses a script that triggers parallel `llvm-cas` invocations on an empty directory.
+// It's intended to ensure there are no racing issues at creation time that will create failures.
+// The last parameter controls how many times it will try.
+// To stress-test this extensively change it locally and pass a large number.
+
+// RUN: %S/Inputs/cas-creation-stress-test.sh llvm-cas %t/cas %s 5