[llvm][cas] Fix thread safety issues with MappedFileRegionBumpPtr

llvm/include/llvm/CAS/MappedFileRegionBumpPtr.h

@@ -50,21 +50,6 @@ class MappedFileRegionBumpPtr {
   create(const Twine &Path, uint64_t Capacity, int64_t BumpPtrOffset,
          function_ref<Error(MappedFileRegionBumpPtr &)> NewFileConstructor);
 
-  /// Create a \c MappedFileRegionBumpPtr., shared across the process via a
-  /// singleton map.
-  ///
-  /// FIXME: Singleton map should be based on sys::fs::UniqueID, but currently
-  /// it is just based on \p Path.
-  ///
-  /// \param Path the path to open the mapped region.
-  /// \param Capacity the maximum size for the mapped file region.
-  /// \param BumpPtrOffset the offset at which to store the bump pointer.
-  /// \param NewFileConstructor is for constructing new files. It has exclusive
-  /// access to the file. Must call \c initializeBumpPtr.
-  static Expected<std::shared_ptr<MappedFileRegionBumpPtr>> createShared(
-      const Twine &Path, uint64_t Capacity, int64_t BumpPtrOffset,
-      function_ref<Error(MappedFileRegionBumpPtr &)> NewFileConstructor);
-
   /// Finish initializing the bump pointer. Must be called by
   /// \c NewFileConstructor.
   void initializeBumpPtr(int64_t BumpPtrOffset);

llvm/lib/CAS/MappedFileRegionBumpPtr.cpp

@@ -52,12 +52,12 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CAS/MappedFileRegionBumpPtr.h"
-#include "llvm/ADT/ScopeExit.h"
+#include "OnDiskCommon.h"
 #include "llvm/ADT/StringMap.h"
-#include <mutex>
 
 using namespace llvm;
 using namespace llvm::cas;
+using namespace llvm::cas::ondisk;
 
 namespace {
 struct FileLockRAII {
@@ -70,7 +70,7 @@ struct FileLockRAII {
   ~FileLockRAII() { consumeError(unlock()); }
 
   Error lock(LockKind LK) {
-    if (std::error_code EC = sys::fs::lockFile(FD, LK == Exclusive))
+    if (std::error_code EC = lockFileThreadSafe(FD, LK == Exclusive))
       return createFileError(Path, EC);
     Locked = LK;
     return Error::success();
@@ -79,7 +79,7 @@ struct FileLockRAII {
   Error unlock() {
     if (Locked) {
       Locked = std::nullopt;
-      if (std::error_code EC = sys::fs::unlockFile(FD))
+      if (std::error_code EC = unlockFileThreadSafe(FD))
         return createFileError(Path, EC);
     }
     return Error::success();
@@ -111,7 +111,8 @@ Expected<MappedFileRegionBumpPtr> MappedFileRegionBumpPtr::create(
 
   // Take shared/reader lock that will be held until we close the file; unlocked
   // by destroyImpl.
-  if (std::error_code EC = sys::fs::lockFile(SharedLockFD, /*Exclusive=*/false))
+  if (std::error_code EC =
+          lockFileThreadSafe(SharedLockFD, /*Exclusive=*/false))
     return createFileError(Path, EC);
 
   // Take shared/reader lock for initialization.
@@ -172,73 +173,22 @@ Expected<MappedFileRegionBumpPtr> MappedFileRegionBumpPtr::create(
   return Result;
 }
 
-Expected<std::shared_ptr<MappedFileRegionBumpPtr>>
-MappedFileRegionBumpPtr::createShared(
-    const Twine &PathTwine, uint64_t Capacity, int64_t BumpPtrOffset,
-    function_ref<Error(MappedFileRegionBumpPtr &)> NewFileConstructor) {
-  struct MapNode {
-    std::mutex Mutex;
-    std::weak_ptr<MappedFileRegionBumpPtr> MFR;
-  };
-  static std::mutex Mutex;
-
-  // FIXME: Map should be by sys::fs::UniqueID instead of by path. Here's how
-  // it should work:
-  //
-  // 1. Open the file.
-  // 2. Stat the file descriptor to get the UniqueID.
-  // 3. Check the map.
-  // 4. If new, pass the open file descriptor to a helper extracted from
-  //    MappedFileRegionBumpPtr::create().
-  static StringMap<MapNode> Regions;
-
-  SmallString<128> PathStorage;
-  const StringRef Path = PathTwine.toStringRef(PathStorage);
-
-  MapNode *Node;
-  {
-    std::lock_guard<std::mutex> Lock(Mutex);
-    Node = &Regions[Path];
-  }
-
-  if (std::shared_ptr<MappedFileRegionBumpPtr> MFR = Node->MFR.lock())
-    return MFR;
-
-  // Construct a new region. Use a fine-grained lock to allow other regions to
-  // be opened concurrently.
-  std::lock_guard<std::mutex> Lock(Node->Mutex);
-
-  // Open / create / initialize files on disk.
-  Expected<MappedFileRegionBumpPtr> ExpectedMFR =
-      MappedFileRegionBumpPtr::create(Path, Capacity, BumpPtrOffset,
-                                      NewFileConstructor);
-  if (!ExpectedMFR)
-    return ExpectedMFR.takeError();
-
-  auto SharedMFR =
-      std::make_shared<MappedFileRegionBumpPtr>(std::move(*ExpectedMFR));
-
-  // Success.
-  Node->MFR = SharedMFR;
-  return std::move(SharedMFR);
-}
-
 void MappedFileRegionBumpPtr::destroyImpl() {
   if (!FD)
     return;
 
   // Drop the shared lock indicating we are no longer accessing the file.
   if (SharedLockFD)
-    (void)sys::fs::unlockFile(*SharedLockFD);
+    (void)unlockFileThreadSafe(*SharedLockFD);
 
   // Attempt to truncate the file if we can get exclusive access. Ignore any
   // errors.
   if (BumpPtr) {
     assert(SharedLockFD && "Must have shared lock file open");
-    if (sys::fs::tryLockFile(*SharedLockFD) == std::error_code()) {
+    if (tryLockFileThreadSafe(*SharedLockFD) == std::error_code()) {
       assert(size() <= capacity());
       (void)sys::fs::resize_file(*FD, size());
-      (void)sys::fs::unlockFile(*SharedLockFD);
+      (void)unlockFileThreadSafe(*SharedLockFD);
     }
   }
 

llvm/lib/CAS/OnDiskCommon.cpp

@@ -12,6 +12,16 @@
 #include "llvm/Support/Process.h"
 #include <mutex>
 #include <optional>
+#include <thread>
+
+#if __has_include(<sys/file.h>)
+#include <sys/file.h>
+#ifdef LOCK_SH
+#define HAVE_FLOCK 1
+#else
+#define HAVE_FLOCK 0
+#endif
+#endif
 
 using namespace llvm;
 
@@ -46,3 +56,55 @@ Expected<std::optional<uint64_t>> cas::ondisk::getOverriddenMaxMappingSize() {
 void cas::ondisk::setMaxMappingSize(uint64_t Size) {
   OnDiskCASMaxMappingSize = Size;
 }
+
+std::error_code cas::ondisk::lockFileThreadSafe(int FD, bool Exclusive) {
+#if HAVE_FLOCK
+  if (flock(FD, Exclusive ? LOCK_EX : LOCK_SH) == 0)
+    return std::error_code();
+  return std::error_code(errno, std::generic_category());
+#elif defined(_WIN32)
+  // On Windows this implementation is thread-safe.
+  return sys::fs::lockFile(FD, Exclusive);
+#else
+  return make_error_code(std::errc::no_lock_available);
+#endif
+}
+
+std::error_code cas::ondisk::unlockFileThreadSafe(int FD) {
+#if HAVE_FLOCK
+  if (flock(FD, LOCK_UN) == 0)
+    return std::error_code();
+  return std::error_code(errno, std::generic_category());
+#elif defined(_WIN32)
+  // On Windows this implementation is thread-safe.
+  return sys::fs::unlockFile(FD);
+#else
+  return make_error_code(std::errc::no_lock_available);
+#endif
+}
+
+std::error_code
+cas::ondisk::tryLockFileThreadSafe(int FD, std::chrono::milliseconds Timeout,
+                                   bool Exclusive) {
+#if HAVE_FLOCK
+  auto Start = std::chrono::steady_clock::now();
+  auto End = Start + Timeout;
+  do {
+    if (flock(FD, (Exclusive ? LOCK_EX : LOCK_SH) | LOCK_NB) == 0)
+      return std::error_code();
+    int Error = errno;
+    if (Error == EWOULDBLOCK) {
+      // Match sys::fs::tryLockFile, which sleeps for 1 ms per attempt.
+      std::this_thread::sleep_for(std::chrono::milliseconds(1));
+      continue;
+    }
+    return std::error_code(Error, std::generic_category());
+  } while (std::chrono::steady_clock::now() < End);
+  return make_error_code(std::errc::no_lock_available);
+#elif defined(_WIN32)
+  // On Windows this implementation is thread-safe.
+  return sys::fs::tryLockFile(FD, Timeout, Exclusive);
+#else
+  return make_error_code(std::errc::no_lock_available);
+#endif
+}

llvm/lib/CAS/OnDiskCommon.h

@@ -10,6 +10,7 @@
 #define LLVM_LIB_CAS_ONDISKCOMMON_H
 
 #include "llvm/Support/Error.h"
+#include <chrono>
 #include <optional>
 
 namespace llvm::cas::ondisk {
@@ -25,6 +26,23 @@ Expected<std::optional<uint64_t>> getOverriddenMaxMappingSize();
 /// created. Set value 0 to use default size.
 void setMaxMappingSize(uint64_t Size);
 
+/// Thread-safe alternative to \c sys::fs::lockFile. This does not support all
+/// the platforms that \c sys::fs::lockFile does, so keep it in the CAS library
+/// for now.
+std::error_code lockFileThreadSafe(int FD, bool Exclusive = true);
+
+/// Thread-safe alternative to \c sys::fs::unlockFile. This does not support all
+/// the platforms that \c sys::fs::lockFile does, so keep it in the CAS library
+/// for now.
+std::error_code unlockFileThreadSafe(int FD);
+
+/// Thread-safe alternative to \c sys::fs::tryLockFile. This does not support
+/// all the platforms that \c sys::fs::lockFile does, so keep it in the CAS
+/// library for now.
+std::error_code tryLockFileThreadSafe(
+    int FD, std::chrono::milliseconds Timeout = std::chrono::milliseconds(0),
+    bool Exclusive = true);
+
 } // namespace llvm::cas::ondisk
 
 #endif // LLVM_LIB_CAS_ONDISKCOMMON_H

llvm/lib/CAS/OnDiskHashMappedTrie.cpp

@@ -138,7 +138,7 @@ class DatabaseFile {
 
 private:
   static Expected<DatabaseFile>
-  get(std::shared_ptr<MappedFileRegionBumpPtr> Alloc) {
+  get(std::unique_ptr<MappedFileRegionBumpPtr> Alloc) {
     if (Error E = validate(Alloc->getRegion()))
       return std::move(E);
     return DatabaseFile(std::move(Alloc));
@@ -148,14 +148,14 @@ class DatabaseFile {
 
   DatabaseFile(MappedFileRegionBumpPtr &Alloc)
       : H(reinterpret_cast<Header *>(Alloc.data())), Alloc(Alloc) {}
-  DatabaseFile(std::shared_ptr<MappedFileRegionBumpPtr> Alloc)
+  DatabaseFile(std::unique_ptr<MappedFileRegionBumpPtr> Alloc)
       : DatabaseFile(*Alloc) {
     OwnedAlloc = std::move(Alloc);
   }
 
   Header *H = nullptr;
   MappedFileRegionBumpPtr &Alloc;
-  std::shared_ptr<MappedFileRegionBumpPtr> OwnedAlloc;
+  std::unique_ptr<MappedFileRegionBumpPtr> OwnedAlloc;
 };
 
 } // end anonymous namespace
@@ -173,14 +173,15 @@ DatabaseFile::create(const Twine &Path, uint64_t Capacity,
   };
 
   // Get or create the file.
-  std::shared_ptr<MappedFileRegionBumpPtr> Alloc;
-  if (Error E = MappedFileRegionBumpPtr::createShared(Path, Capacity,
-                                                      offsetof(Header, BumpPtr),
-                                                      NewFileConstructor)
+  MappedFileRegionBumpPtr Alloc;
+  if (Error E = MappedFileRegionBumpPtr::create(Path, Capacity,
+                                                offsetof(Header, BumpPtr),
+                                                NewFileConstructor)
                     .moveInto(Alloc))
     return std::move(E);
 
-  return DatabaseFile::get(std::move(Alloc));
+  return DatabaseFile::get(
+      std::make_unique<MappedFileRegionBumpPtr>(std::move(Alloc)));
 }
 
 void DatabaseFile::addTable(TableHandle Table) {

llvm/lib/CAS/UnifiedOnDiskCache.cpp

@@ -50,6 +50,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CAS/UnifiedOnDiskCache.h"
+#include "OnDiskCommon.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/CAS/OnDiskKeyValueDB.h"
 #include "llvm/Support/Errc.h"
@@ -182,7 +183,7 @@ UnifiedOnDiskCache::open(StringRef RootPath, std::optional<uint64_t> SizeLimit,
   // from creating a new chain (essentially while a \p UnifiedOnDiskCache
   // instance holds a shared lock the storage for the primary directory will
   // grow unrestricted).
-  if (std::error_code EC = sys::fs::lockFile(LockFD, /*Exclusive=*/false))
+  if (std::error_code EC = lockFileThreadSafe(LockFD, /*Exclusive=*/false))
     return createFileError(PathBuf, EC);
 
   SmallVector<std::string, 4> DBDirs;
@@ -295,7 +296,7 @@ Error UnifiedOnDiskCache::close(bool CheckSizeLimit) {
   UpstreamKVDB.reset();
   PrimaryGraphDB.reset();
   UpstreamGraphDB = nullptr;
-  if (std::error_code EC = sys::fs::unlockFile(LockFD))
+  if (std::error_code EC = unlockFileThreadSafe(LockFD))
     return createFileError(RootPath, EC);
 
   if (!ExceededSizeLimit)
@@ -305,13 +306,13 @@ Error UnifiedOnDiskCache::close(bool CheckSizeLimit) {
   // exclusive lock in order to create a new primary directory for next time
   // this \p UnifiedOnDiskCache path is opened.
 
-  if (std::error_code EC = sys::fs::tryLockFile(
+  if (std::error_code EC = tryLockFileThreadSafe(
           LockFD, std::chrono::milliseconds(0), /*Exclusive=*/true)) {
     if (EC == errc::no_lock_available)
       return Error::success(); // couldn't get exclusive lock, give up.
     return createFileError(RootPath, EC);
   }
-  auto _2 = make_scope_exit([&]() { sys::fs::unlockFile(LockFD); });
+  auto _2 = make_scope_exit([&]() { unlockFileThreadSafe(LockFD); });
 
   // Managed to get an exclusive lock which means there are no other open
   // \p UnifiedOnDiskCache instances for the same path, so we can safely start a

llvm/test/check-lock-files.ll

@@ -0,0 +1,6 @@
+# REQUIRES: ondisk_cas
+
+# Multi-threaded test that CAS lock files protecting the shared data are working.
+
+# RUN: rm -rf %t/cas
+# RUN: llvm-cas -cas %t/cas -check-lock-files