[CAS] Improve MappedFileRegionBumpPtr

Improve MappedFileRegionBumpPtr so it can handle being opened with
different capacities.

Author: cachemeifyoucan

llvm/include/llvm/CAS/MappedFileRegionBumpPtr.h

@@ -5,11 +5,16 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+//
+/// \file
+/// This file declares interface for MappedFileRegionBumpPtr, a bump pointer
+/// allocator, backed by a memory-mapped file.
+///
+//===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CAS_MAPPEDFILEREGIONBUMPPTR_H
 #define LLVM_CAS_MAPPEDFILEREGIONBUMPPTR_H
 
-#include "llvm/Config/llvm-config.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/FileSystem.h"
 #include <atomic>
@@ -18,7 +23,7 @@ namespace llvm::cas {
 
 namespace ondisk {
 class OnDiskCASLogger;
-}
+} // namespace ondisk
 
 /// Allocator for an owned mapped file region that supports thread-safe and
 /// process-safe bump pointer allocation.
@@ -36,28 +41,34 @@ class OnDiskCASLogger;
 /// in the same process since file locks will misbehave. Clients should
 /// coordinate (somehow).
 ///
-/// \note Currently we allocate the whole file without sparseness on Windows.
-///
 /// Provides 8-byte alignment for all allocations.
 class MappedFileRegionBumpPtr {
 public:
   using RegionT = sys::fs::mapped_file_region;
 
+  /// Header for MappedFileRegionBumpPtr. It can be configured to be located
+  /// at any location within the file and the allocation will be appended after
+  /// the header.
+  struct Header {
+    std::atomic<uint64_t> BumpPtr;
+    std::atomic<uint64_t> AllocatedSize;
+  };
+
   /// Create a \c MappedFileRegionBumpPtr.
   ///
   /// \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 HeaderOffset the offset at which to store the header. This is so
+  /// that information can be stored before the header, like a file magic.
   /// \param NewFileConstructor is for constructing new files. It has exclusive
   /// access to the file. Must call \c initializeBumpPtr.
   static Expected<MappedFileRegionBumpPtr>
-  create(const Twine &Path, uint64_t Capacity, int64_t BumpPtrOffset,
+  create(const Twine &Path, uint64_t Capacity, uint64_t HeaderOffset,
          std::shared_ptr<ondisk::OnDiskCASLogger> Logger,
          function_ref<Error(MappedFileRegionBumpPtr &)> NewFileConstructor);
 
-  /// Finish initializing the bump pointer. Must be called by
-  /// \c NewFileConstructor.
-  void initializeBumpPtr(int64_t BumpPtrOffset);
+  /// Finish initializing the header. Must be called by \c NewFileConstructor.
+  void initializeHeader(uint64_t HeaderOffset);
 
   /// Minimum alignment for allocations, currently hardcoded to 8B.
   static constexpr Align getAlign() {
@@ -108,14 +119,12 @@ class MappedFileRegionBumpPtr {
   }
 
 private:
-  struct Header {
-    std::atomic<int64_t> BumpPtr;
-    std::atomic<int64_t> AllocatedSize;
-  };
   RegionT Region;
   Header *H = nullptr;
   std::string Path;
+  // File descriptor for the main storage file.
   std::optional<int> FD;
+  // File descriptor for the file used as reader/writer lock.
   std::optional<int> SharedLockFD;
   std::shared_ptr<ondisk::OnDiskCASLogger> Logger = nullptr;
 };

llvm/lib/CAS/MappedFileRegionBumpPtr.cpp

@@ -1,11 +1,11 @@
-//===- MappedFileRegionBumpPtr.cpp ------------------------------------===//
+//===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
-/// \file
+/// \file Implements MappedFileRegionBumpPtr.
 ///
 /// A bump pointer allocator, backed by a memory-mapped file.
 ///
@@ -20,12 +20,14 @@
 /// and across multiple processes without locking for every read. Our current
 /// implementation strategy is:
 ///
-/// 1. Use \c ftruncate (\c sys::fs::resize_file) to grow the file to its max
-///    size (typically several GB). Many modern filesystems will create a sparse
-///    file, so that the trailing unused pages do not take space on disk.
-/// 2. Call \c mmap (\c sys::fs::mapped_file_region)
+/// 1. Use \c sys::fs::resize_file_sparse to grow the file to its max size
+///    (typically several GB). If the file system doesn't support sparse file,
+///    this may return a fully allocated file.
+/// 2. Call \c sys::fs::mapped_file_region to map the entire file.
 /// 3. [Automatic as part of 2.]
-/// 4. [Automatic as part of 2.]
+/// 4. If supported, use \c fallocate or similiar APIs to ensure the file system
+///    storage for the sparse file so we won't end up with partial file if the
+///    disk is out of space.
 ///
 /// Additionally, we attempt to resize the file to its actual data size when
 /// closing the mapping, if this is the only concurrent instance. This is done
@@ -35,10 +37,10 @@
 /// which typically loses sparseness. These mitigations only work while the file
 /// is not in use.
 ///
-/// FIXME: we assume that all concurrent users of the file will use the same
-/// value for Capacity. Otherwise a process with a larger capacity can write
-/// data that is "out of bounds" for processes with smaller capacity. Currently
-/// this is true in the CAS.
+/// If different values of the capacity is used for concurrent users of the same
+/// mapping, the capacity is determined by the first value used to open the
+/// file. It is a requirement for the users to always open the file with the
+/// same \c HeaderOffset, otherwise the behavior is undefined.
 ///
 /// To support resizing, we use two separate file locks:
 /// 1. We use a shared reader lock on a ".shared" file until destruction.
@@ -54,7 +56,6 @@
 #include "llvm/CAS/MappedFileRegionBumpPtr.h"
 #include "OnDiskCommon.h"
 #include "llvm/CAS/OnDiskCASLogger.h"
-#include "llvm/Support/Compiler.h"
 
 #if LLVM_ON_UNIX
 #include <sys/stat.h>
@@ -82,6 +83,10 @@ struct FileLockRAII {
   ~FileLockRAII() { consumeError(unlock()); }
 
   Error lock(sys::fs::LockKind LK) {
+    // Try unlock first. If not locked, this is no-op.
+    if (auto E = unlock())
+      return E;
+
     if (std::error_code EC = lockFileThreadSafe(FD, LK))
       return createFileError(Path, EC);
     Locked = LK;
@@ -107,9 +112,15 @@ struct FileSizeInfo {
 } // end anonymous namespace
 
 Expected<MappedFileRegionBumpPtr> MappedFileRegionBumpPtr::create(
-    const Twine &Path, uint64_t Capacity, int64_t BumpPtrOffset,
+    const Twine &Path, uint64_t Capacity, uint64_t HeaderOffset,
     std::shared_ptr<ondisk::OnDiskCASLogger> Logger,
     function_ref<Error(MappedFileRegionBumpPtr &)> NewFileConstructor) {
+  uint64_t MinCapacity = HeaderOffset + sizeof(Header);
+  if (Capacity < MinCapacity)
+    return createStringError(
+        std::make_error_code(std::errc::invalid_argument),
+        "capacity is too small to hold MappedFileRegionBumpPtr");
+
   MappedFileRegionBumpPtr Result;
   Result.Path = Path.str();
   Result.Logger = std::move(Logger);
@@ -146,66 +157,82 @@ Expected<MappedFileRegionBumpPtr> MappedFileRegionBumpPtr::create(
   if (!FileSize)
     return createFileError(Result.Path, FileSize.getError());
 
+  // If the size is smaller than the capacity, we need to initialize the file.
+  // It maybe empty, or may have been shrunk during a previous close.
   if (FileSize->Size < Capacity) {
     // Lock the file exclusively so only one process will do the initialization.
-    if (Error E = InitLock.unlock())
-      return std::move(E);
     if (Error E = InitLock.lock(sys::fs::LockKind::Exclusive))
       return std::move(E);
     // Retrieve the current size now that we have exclusive access.
     FileSize = FileSizeInfo::get(File);
     if (!FileSize)
-        return createFileError(Result.Path, FileSize.getError());
+      return createFileError(Result.Path, FileSize.getError());
   }
 
-  // At this point either the file is still under-sized, or we have the size for
-  // the completely initialized file.
-
-  if (FileSize->Size < Capacity) {
-    // We are initializing the file; it may be empty, or may have been shrunk
-    // during a previous close.
-    // FIXME: Detect a case where someone opened it with a smaller capacity.
+  uint64_t MappingSize = FileSize->Size;
+  // If the size is still smaller than the minimal required size, we need to
+  // resize the file to the capacity.
+  if (FileSize->Size < MinCapacity) {
     assert(InitLock.Locked == sys::fs::LockKind::Exclusive);
     if (std::error_code EC = sys::fs::resize_file_sparse(FD, Capacity))
       return createFileError(Result.Path, EC);
-
     if (Result.Logger)
       Result.Logger->log_MappedFileRegionBumpPtr_resizeFile(
           Result.Path, FileSize->Size, Capacity);
-  } else {
-    // Someone else initialized it.
-    Capacity = FileSize->Size;
+    MappingSize = Capacity;
   }
 
   // Create the mapped region.
   {
     std::error_code EC;
     sys::fs::mapped_file_region Map(
-        File, sys::fs::mapped_file_region::readwrite, Capacity, 0, EC);
+        File, sys::fs::mapped_file_region::readwrite, MappingSize, 0, EC);
     if (EC)
       return createFileError(Result.Path, EC);
     Result.Region = std::move(Map);
   }
 
-  if (FileSize->Size == 0) {
+  if (FileSize->Size < MinCapacity) {
     assert(InitLock.Locked == sys::fs::LockKind::Exclusive);
-    // We are creating a new file; run the constructor.
+    // If we need to fully initialize the file, call NewFileConstructor.
     if (Error E = NewFileConstructor(Result))
       return std::move(E);
-  } else {
-    Result.initializeBumpPtr(BumpPtrOffset);
+  } else
+    Result.initializeHeader(HeaderOffset);
+
+  if (Result.H->BumpPtr >= FileSize->Size && FileSize->Size < Capacity) {
+    assert(InitLock.Locked == sys::fs::LockKind::Exclusive);
+    // If the BumpPtr larger than or equal to the size of the file (it can be
+    // larger if process is terminated when the out of memory allocation
+    // happens) and smaller than capacity, this was shrunken by a previous
+    // close, resize back to capacity and re-initialize the mapped_file_region.
+    Result.Region.unmap();
+    if (std::error_code EC = sys::fs::resize_file_sparse(FD, Capacity))
+      return createFileError(Result.Path, EC);
+    if (Result.Logger)
+      Result.Logger->log_MappedFileRegionBumpPtr_resizeFile(
+          Result.Path, FileSize->Size, Capacity);
+
+    std::error_code EC;
+    sys::fs::mapped_file_region Map(
+        File, sys::fs::mapped_file_region::readwrite, Capacity, 0, EC);
+    if (EC)
+      return createFileError(Result.Path, EC);
+    Result.Region = std::move(Map);
+    Result.initializeHeader(HeaderOffset);
   }
 
-  if (FileSize->Size < Capacity && FileSize->AllocatedSize < Capacity) {
-    // We are initializing the file; sync the allocated size in case it
-    // changed when truncating or during construction.
+  if (InitLock.Locked == sys::fs::LockKind::Exclusive) {
+    // If holding an exclusive lock, we might have resized the file and
+    // performed some read/write to the file. Query the file size again to make
+    // sure everything is up-to-date. Otherwise, FileSize info is already
+    // up-to-date.
     FileSize = FileSizeInfo::get(File);
     if (!FileSize)
       return createFileError(Result.Path, FileSize.getError());
-    assert(InitLock.Locked == sys::fs::LockKind::Exclusive);
-    Result.H->AllocatedSize.exchange(FileSize->AllocatedSize);
   }
 
+  Result.H->AllocatedSize.exchange(FileSize->AllocatedSize);
   return Result;
 }
 
@@ -226,10 +253,11 @@ void MappedFileRegionBumpPtr::destroyImpl() {
       size_t Capacity = capacity();
       // sync to file system to make sure all contents are up-to-date.
       (void)Region.sync();
+      // unmap the file before resizing since that is the requirement for
+      // some platforms.
       Region.unmap();
       (void)sys::fs::resize_file(*FD, Size);
       (void)unlockFileThreadSafe(*SharedLockFD);
-
       if (Logger)
         Logger->log_MappedFileRegionBumpPtr_resizeFile(Path, Capacity, Size);
     }
@@ -251,20 +279,19 @@ void MappedFileRegionBumpPtr::destroyImpl() {
     Logger->log_MappedFileRegionBumpPtr_close(Path);
 }
 
-void MappedFileRegionBumpPtr::initializeBumpPtr(int64_t BumpPtrOffset) {
+void MappedFileRegionBumpPtr::initializeHeader(uint64_t HeaderOffset) {
   assert(capacity() < (uint64_t)INT64_MAX && "capacity must fit in int64_t");
-  int64_t BumpPtrEndOffset = BumpPtrOffset + sizeof(decltype(*H));
-  assert(BumpPtrEndOffset <= (int64_t)capacity() &&
+  uint64_t HeaderEndOffset = HeaderOffset + sizeof(decltype(*H));
+  assert(HeaderEndOffset <= capacity() &&
          "Expected end offset to be pre-allocated");
-  assert(isAligned(Align::Of<decltype(*H)>(), BumpPtrOffset) &&
+  assert(isAligned(Align::Of<decltype(*H)>(), HeaderOffset) &&
          "Expected end offset to be aligned");
-  H = reinterpret_cast<decltype(H)>(data() + BumpPtrOffset);
-
-  int64_t ExistingValue = 0;
-  if (!H->BumpPtr.compare_exchange_strong(ExistingValue, BumpPtrEndOffset))
-    assert(ExistingValue >= BumpPtrEndOffset &&
-           "Expected 0, or past the end of the BumpPtr itself");
+  H = reinterpret_cast<decltype(H)>(data() + HeaderOffset);
 
+  uint64_t ExistingValue = 0;
+  if (!H->BumpPtr.compare_exchange_strong(ExistingValue, HeaderEndOffset))
+    assert(ExistingValue >= HeaderEndOffset &&
+           "Expected 0, or past the end of the header itself");
   if (Logger)
     Logger->log_MappedFileRegionBumpPtr_create(Path, *FD, data(), capacity(),
                                                size());
@@ -277,16 +304,16 @@ static Error createAllocatorOutOfSpaceError() {
 
 Expected<int64_t> MappedFileRegionBumpPtr::allocateOffset(uint64_t AllocSize) {
   AllocSize = alignTo(AllocSize, getAlign());
-  int64_t OldEnd = H->BumpPtr.fetch_add(AllocSize);
-  int64_t NewEnd = OldEnd + AllocSize;
-  if (LLVM_UNLIKELY(NewEnd > (int64_t)capacity())) {
+  uint64_t OldEnd = H->BumpPtr.fetch_add(AllocSize);
+  uint64_t NewEnd = OldEnd + AllocSize;
+  if (LLVM_UNLIKELY(NewEnd > capacity())) {
     // Return the allocation. If the start already passed the end, that means
     // some other concurrent allocations already consumed all the capacity.
     // There is no need to return the original value. If the start was not
     // passed the end, current allocation certainly bumped it passed the end.
     // All other allocation afterwards must have failed and current allocation
     // is in charge of return the allocation back to a valid value.
-    if (OldEnd <= (int64_t)capacity())
+    if (OldEnd <= capacity())
       (void)H->BumpPtr.exchange(OldEnd);
 
     if (Logger)
@@ -296,12 +323,13 @@ Expected<int64_t> MappedFileRegionBumpPtr::allocateOffset(uint64_t AllocSize) {
     return createAllocatorOutOfSpaceError();
   }
 
-  int64_t DiskSize = H->AllocatedSize;
+  uint64_t DiskSize = H->AllocatedSize;
   if (LLVM_UNLIKELY(NewEnd > DiskSize)) {
-    int64_t NewSize;
+    uint64_t NewSize;
     // The minimum increment is a page, but allocate more to amortize the cost.
-    constexpr int64_t Increment = 1 * 1024 * 1024; // 1 MB
-    if (Error E = preallocateFileTail(*FD, DiskSize, DiskSize + Increment).moveInto(NewSize))
+    constexpr uint64_t Increment = 1 * 1024 * 1024; // 1 MB
+    if (Error E = preallocateFileTail(*FD, DiskSize, DiskSize + Increment)
+                      .moveInto(NewSize))
       return std::move(E);
     assert(NewSize >= DiskSize + Increment);
     // FIXME: on Darwin this can under-count the size if there is a race to

llvm/lib/CAS/OnDiskHashMappedTrie.cpp

@@ -123,7 +123,7 @@ class DatabaseFile {
     uint64_t Magic;
     uint64_t Version;
     std::atomic<int64_t> RootTableOffset;
-    std::atomic<int64_t> BumpPtr;
+    MappedFileRegionBumpPtr::Header MappedFileHeader;
   };
 
   const Header &getHeader() { return *H; }
@@ -185,16 +185,16 @@ DatabaseFile::create(const Twine &Path, uint64_t Capacity,
       return createTableConfigError(std::errc::argument_out_of_domain,
                                     Path.str(), "datafile",
                                     "Allocator too small for header");
-    (void)new (Alloc.data()) Header{getMagic(), getVersion(), {0}, {0}};
-    Alloc.initializeBumpPtr(offsetof(Header, BumpPtr));
+    (void)new (Alloc.data()) Header{getMagic(), getVersion(), {0}, {}};
+    Alloc.initializeHeader(offsetof(Header, MappedFileHeader));
     DatabaseFile DB(Alloc);
     return NewDBConstructor(DB);
   };
 
   // Get or create the file.
   MappedFileRegionBumpPtr Alloc;
   if (Error E = MappedFileRegionBumpPtr::create(
-                    Path, Capacity, offsetof(Header, BumpPtr),
+                    Path, Capacity, offsetof(Header, MappedFileHeader),
                     std::move(Logger), NewFileConstructor)
                     .moveInto(Alloc))
     return std::move(E);
@@ -264,7 +264,7 @@ Error DatabaseFile::validate(MappedFileRegion &Region) {
                              "database: wrong version");
 
   // Check the bump-ptr, which should point past the header.
-  if (H->BumpPtr.load() < (int64_t)sizeof(Header))
+  if (H->MappedFileHeader.BumpPtr.load() < (int64_t)sizeof(Header))
     return createStringError(std::errc::invalid_argument,
                              "database: corrupt bump-ptr");