[CAS] Add MappedFileRegionArena (#114099)

Add MappedFileRegionArena which can be served as a file system backed
persistent memory allocator. The allocator works like a
BumpPtrAllocator,
and is designed to be thread safe and process safe.

The implementation relies on the POSIX compliance of file system and
doesn't work on all file systems. If the file system supports lazy tail
(doesn't allocate disk space if the tail of the large file is not used),
user has more flexibility to declare a larger capacity.

The allocator works by using a atomically updated bump ptr at a location
that can be customized by the user. The atomic pointer points to the
next available space to allocate, and the allocator will resize/truncate
to current usage once all clients closed the allocator.

Windows implementation contributed by: @hjyamauchi

Author: cachemeifyoucan

llvm/CMakeLists.txt

@@ -878,6 +878,7 @@ option (LLVM_ENABLE_SPHINX "Use Sphinx to generate llvm documentation." OFF)
 option (LLVM_ENABLE_OCAMLDOC "Build OCaml bindings documentation." ON)
 option (LLVM_ENABLE_BINDINGS "Build bindings." ON)
 option (LLVM_ENABLE_TELEMETRY "Enable the telemetry library. If set to OFF, library cannot be enabled after build (eg., at runtime)" ON)
+option (LLVM_ENABLE_ONDISK_CAS "Build OnDiskCAS." ON)
 
 set(LLVM_INSTALL_DOXYGEN_HTML_DIR "${CMAKE_INSTALL_DOCDIR}/llvm/doxygen-html"
     CACHE STRING "Doxygen-generated HTML documentation install directory")

llvm/include/llvm/CAS/MappedFileRegionArena.h

@@ -0,0 +1,130 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+/// This file declares interface for MappedFileRegionArena, a bump pointer
+/// allocator, backed by a memory-mapped file.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CAS_MAPPEDFILEREGIONARENA_H
+#define LLVM_CAS_MAPPEDFILEREGIONARENA_H
+
+#include "llvm/Support/Alignment.h"
+#include "llvm/Support/FileSystem.h"
+#include <atomic>
+
+namespace llvm::cas {
+
+/// Allocator for an owned mapped file region that supports thread-safe and
+/// process-safe bump pointer allocation.
+///
+/// This allocator is designed to create a sparse file when supported by the
+/// filesystem's \c ftruncate so that it can be used with a large maximum size.
+/// It will also attempt to shrink the underlying file down to its current
+/// allocation size when the last concurrent mapping is closed.
+///
+/// Process-safe. Uses file locks when resizing the file during initialization
+/// and destruction.
+///
+/// Thread-safe. Requires OS support thread-safe file lock.
+///
+/// Provides 8-byte alignment for all allocations.
+class MappedFileRegionArena {
+public:
+  using RegionT = sys::fs::mapped_file_region;
+
+  /// Header for MappedFileRegionArena. It can be configured to be located
+  /// at any location within the file and the allocation will be appended after
+  /// the header.
+  struct Header {
+    // BumpPtr for new allocation.
+    std::atomic<uint64_t> BumpPtr;
+    // Allocated size on disk.
+    std::atomic<uint64_t> AllocatedSize;
+    // Capacity of the file.
+    std::atomic<uint64_t> Capacity;
+    // Offset from the beginning of the file to this header (for verification).
+    std::atomic<uint64_t> HeaderOffset;
+  };
+
+  /// Create a \c MappedFileRegionArena.
+  ///
+  /// \param Path the path to open the mapped region.
+  /// \param Capacity the maximum size for the mapped file region.
+  /// \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<MappedFileRegionArena>
+  create(const Twine &Path, uint64_t Capacity, uint64_t HeaderOffset,
+         function_ref<Error(MappedFileRegionArena &)> NewFileConstructor);
+
+  /// Minimum alignment for allocations, currently hardcoded to 8B.
+  static constexpr Align getAlign() {
+    // Trick Align into giving us '8' as a constexpr.
+    struct alignas(8) T {};
+    static_assert(alignof(T) == 8, "Tautology failed?");
+    return Align::Of<T>();
+  }
+
+  /// Allocate at least \p AllocSize. Rounds up to \a getAlign().
+  Expected<char *> allocate(uint64_t AllocSize) {
+    auto Offset = allocateOffset(AllocSize);
+    if (LLVM_UNLIKELY(!Offset))
+      return Offset.takeError();
+    return data() + *Offset;
+  }
+  /// Allocate, returning the offset from \a data() instead of a pointer.
+  Expected<int64_t> allocateOffset(uint64_t AllocSize);
+
+  char *data() const { return Region.data(); }
+  uint64_t size() const { return H->BumpPtr; }
+  uint64_t capacity() const { return Region.size(); }
+
+  RegionT &getRegion() { return Region; }
+
+  ~MappedFileRegionArena() { destroyImpl(); }
+
+  MappedFileRegionArena() = default;
+  MappedFileRegionArena(MappedFileRegionArena &&RHS) { moveImpl(RHS); }
+  MappedFileRegionArena &operator=(MappedFileRegionArena &&RHS) {
+    destroyImpl();
+    moveImpl(RHS);
+    return *this;
+  }
+
+  MappedFileRegionArena(const MappedFileRegionArena &) = delete;
+  MappedFileRegionArena &operator=(const MappedFileRegionArena &) = delete;
+
+private:
+  // initialize header from offset.
+  void initializeHeader(uint64_t HeaderOffset);
+
+  void destroyImpl();
+  void moveImpl(MappedFileRegionArena &RHS) {
+    std::swap(Region, RHS.Region);
+    std::swap(H, RHS.H);
+    std::swap(Path, RHS.Path);
+    std::swap(FD, RHS.FD);
+    std::swap(SharedLockFD, RHS.SharedLockFD);
+  }
+
+private:
+  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;
+};
+
+} // namespace llvm::cas
+
+#endif // LLVM_CAS_MAPPEDFILEREGIONARENA_H

llvm/include/llvm/Config/llvm-config.h.cmake

@@ -146,4 +146,7 @@
    coverage bugs, and to 0 otherwise. */
 #cmakedefine01 LLVM_ENABLE_DEBUGLOC_TRACKING_ORIGIN
 
+/* Define to 1 to enable LLVM OnDisk Content Addressable Storage */
+#cmakedefine01 LLVM_ENABLE_ONDISK_CAS
+
 #endif

llvm/include/llvm/Support/FileSystem.h

@@ -410,6 +410,11 @@ LLVM_ABI std::error_code copy_file(const Twine &From, int ToFD);
 ///          platform-specific error_code.
 LLVM_ABI std::error_code resize_file(int FD, uint64_t Size);
 
+/// Resize path to size with sparse files explicitly enabled. It uses
+/// FSCTL_SET_SPARSE On Windows. This is the same as resize_file on
+/// non-Windows
+LLVM_ABI std::error_code resize_file_sparse(int FD, uint64_t Size);
+
 /// Resize \p FD to \p Size before mapping \a mapped_file_region::readwrite. On
 /// non-Windows, this calls \a resize_file(). On Windows, this is a no-op,
 /// since the subsequent mapping (via \c CreateFileMapping) automatically

llvm/lib/CAS/CMakeLists.txt

@@ -3,7 +3,9 @@ add_llvm_component_library(LLVMCAS
   ActionCaches.cpp
   BuiltinCAS.cpp
   InMemoryCAS.cpp
+  MappedFileRegionArena.cpp
   ObjectStore.cpp
+  OnDiskCommon.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/CAS