Handle larger OS page sizes in VM eviction

src/include/storage/buffer_manager/buffer_manager.h

@@ -275,12 +275,14 @@ class BufferManager {
 
     uint64_t freeUsedMemory(uint64_t size);
 
-    void releaseFrameForPage(FileHandle& fileHandle [[maybe_unused]],
+    uint64_t releaseFrameForPage(FileHandle& fileHandle [[maybe_unused]],
         common::page_idx_t pageIdx [[maybe_unused]]) {
 #if BM_MALLOC
         // Page is freed instead in PageState::resetToEvicted
+        return fileHandle.getPageSize();
 #else
-        vmRegions[fileHandle.getPageSizeClass()]->releaseFrame(fileHandle.getFrameIdx(pageIdx));
+        return vmRegions[fileHandle.getPageSizeClass()]->releaseFrame(
+            fileHandle.getFrameIdx(pageIdx));
 #endif
     }
 

src/include/storage/buffer_manager/vm_region.h

@@ -1,6 +1,9 @@
 #pragma once
 
+#include <atomic>
+#include <memory>
 #include <mutex>
+#include <vector>
 
 #include "common/constants.h"
 #include "common/types/types.h"
@@ -22,8 +25,13 @@ class VMRegion {
 
     common::frame_group_idx_t addNewFrameGroup();
 
+    // Mark the frame as resident. Returns the number of bytes that should be charged to the
+    // buffer pool for this frame.
+    uint64_t claimFrame(common::frame_idx_t frameIdx);
+
     // Use `MADV_DONTNEED` to release physical memory associated with this frame.
-    void releaseFrame(common::frame_idx_t frameIdx) const;
+    // Returns the number of bytes that should be released from the buffer pool accounting.
+    uint64_t releaseFrame(common::frame_idx_t frameIdx);
 
     // Returns true if the memory address is within the reserved virtual memory region
     bool contains(const uint8_t* address) const {
@@ -37,13 +45,21 @@ class VMRegion {
     inline uint64_t getMaxRegionSize() const {
         return maxNumFrameGroups * frameSize * common::StorageConstants::PAGE_GROUP_SIZE;
     }
+    inline uint64_t getFrameGroupSize() const {
+        return static_cast<uint64_t>(frameSize) * common::StorageConstants::PAGE_GROUP_SIZE;
+    }
+    uint64_t getDiscardGranuleIdxInFrameGroup(common::frame_idx_t frameIdx) const;
+    uint64_t getFrameGroupIdx(common::frame_idx_t frameIdx) const;
 
 private:
     std::mutex mtx;
     uint8_t* region;
     uint32_t frameSize;
+    uint64_t discardGranuleSize;
+    uint64_t numDiscardGranulesPerFrameGroup;
     uint64_t numFrameGroups;
     uint64_t maxNumFrameGroups;
+    std::vector<std::unique_ptr<std::atomic<uint16_t>[]>> residentFramesPerDiscardGranule;
 };
 
 } // namespace storage

src/storage/buffer_manager/buffer_manager.cpp

@@ -343,8 +343,15 @@ void BufferManager::removeEvictedCandidates() {
 // and return false, otherwise, we load the page to its corresponding frame and return true.
 bool BufferManager::claimAFrame(FileHandle& fileHandle, page_idx_t pageIdx,
     PageReadPolicy pageReadPolicy) {
-    page_offset_t pageSizeToClaim = fileHandle.getPageSize();
+    uint64_t pageSizeToClaim = fileHandle.getPageSize();
+#if !BM_MALLOC
+    pageSizeToClaim =
+        vmRegions[fileHandle.getPageSizeClass()]->claimFrame(fileHandle.getFrameIdx(pageIdx));
+#endif
     if (!reserve(pageSizeToClaim)) {
+#if !BM_MALLOC
+        vmRegions[fileHandle.getPageSizeClass()]->releaseFrame(fileHandle.getFrameIdx(pageIdx));
+#endif
         return false;
     }
 #if _WIN32 && !BM_MALLOC
@@ -453,8 +460,7 @@ uint64_t BufferManager::tryEvictPage(std::atomic<EvictionCandidate>& _candidate)
     // is dirty. Finally remove the page from the frame and reset the page to EVICTED.
     auto& fileHandle = *fileHandles[candidate.fileIdx];
     fileHandle.flushPageIfDirtyWithoutLock(candidate.pageIdx);
-    auto numBytesFreed = fileHandle.getPageSize();
-    releaseFrameForPage(fileHandle, candidate.pageIdx);
+    auto numBytesFreed = releaseFrameForPage(fileHandle, candidate.pageIdx);
     evictionQueue.clear(_candidate);
     pageState.resetToEvicted();
     return numBytesFreed;
@@ -533,8 +539,8 @@ void BufferManager::removePageFromFrame(FileHandle& fileHandle, page_idx_t pageI
     if (shouldFlush) {
         fileHandle.flushPageIfDirtyWithoutLock(pageIdx);
     }
-    releaseFrameForPage(fileHandle, pageIdx);
-    freeUsedMemory(fileHandle.getPageSize());
+    const auto numBytesFreed = releaseFrameForPage(fileHandle, pageIdx);
+    freeUsedMemory(numBytesFreed);
     pageState->resetToEvicted();
 }
 

src/storage/buffer_manager/vm_region.cpp

@@ -1,18 +1,21 @@
 #include "storage/buffer_manager/vm_region.h"
 
+#include <algorithm>
+
 #include "common/system_config.h"
 #include "common/system_message.h"
+#include <format>
 
 #ifdef _WIN32
 #include <errhandlingapi.h>
 #include <handleapi.h>
 #include <memoryapi.h>
 #else
 #include <sys/mman.h>
-
-#include <format>
+#include <unistd.h>
 #endif
 
+#include "common/assert.h"
 #include "common/exception/buffer_manager.h"
 
 #ifndef MAP_NORESERVE
@@ -29,6 +32,21 @@ VMRegion::VMRegion(PageSizeClass pageSizeClass, uint64_t maxRegionSize) : numFra
         throw BufferManagerException("maxRegionSize is beyond the max available mmap region size.");
     }
     frameSize = pageSizeClass == REGULAR_PAGE ? LBUG_PAGE_SIZE : TEMP_PAGE_SIZE;
+    discardGranuleSize = frameSize;
+#ifndef _WIN32
+    const auto osPageSize = sysconf(_SC_PAGESIZE);
+    if (osPageSize <= 0) {
+        throw BufferManagerException("Failed to detect the operating system page size.");
+    }
+    discardGranuleSize = std::max<uint64_t>(frameSize, static_cast<uint64_t>(osPageSize));
+#endif
+    if (discardGranuleSize % frameSize != 0 || getFrameGroupSize() % discardGranuleSize != 0) {
+        throw BufferManagerException(std::format(
+            "Unsupported page size combination: frame size {}, discard granule size {}, frame "
+            "group size {}.",
+            frameSize, discardGranuleSize, getFrameGroupSize()));
+    }
+    numDiscardGranulesPerFrameGroup = getFrameGroupSize() / discardGranuleSize;
     const auto numBytesForFrameGroup = frameSize * StorageConstants::PAGE_GROUP_SIZE;
     maxNumFrameGroups = (maxRegionSize + numBytesForFrameGroup - 1) / numBytesForFrameGroup;
 #ifdef _WIN32
@@ -58,20 +76,41 @@ VMRegion::~VMRegion() {
 #endif
 }
 
-void VMRegion::releaseFrame(frame_idx_t frameIdx) const {
+uint64_t VMRegion::claimFrame(frame_idx_t frameIdx) {
+    const auto frameGroupIdx = getFrameGroupIdx(frameIdx);
+    DASSERT(frameGroupIdx < residentFramesPerDiscardGranule.size());
+    const auto granuleIdx = getDiscardGranuleIdxInFrameGroup(frameIdx);
+    auto& numResidentFrames = residentFramesPerDiscardGranule[frameGroupIdx][granuleIdx];
+    const auto previousNumResidentFrames = numResidentFrames.fetch_add(1);
+    DASSERT(previousNumResidentFrames < discardGranuleSize / frameSize);
+    return previousNumResidentFrames == 0 ? discardGranuleSize : 0;
+}
+
+uint64_t VMRegion::releaseFrame(frame_idx_t frameIdx) {
+    const auto frameGroupIdx = getFrameGroupIdx(frameIdx);
+    DASSERT(frameGroupIdx < residentFramesPerDiscardGranule.size());
+    const auto granuleIdx = getDiscardGranuleIdxInFrameGroup(frameIdx);
+    auto& numResidentFrames = residentFramesPerDiscardGranule[frameGroupIdx][granuleIdx];
+    const auto previousNumResidentFrames = numResidentFrames.fetch_sub(1);
+    DASSERT(previousNumResidentFrames > 0);
+    if (previousNumResidentFrames != 1) {
+        return 0;
+    }
+    auto* discardGranule =
+        region + frameGroupIdx * getFrameGroupSize() + granuleIdx * discardGranuleSize;
 #ifdef _WIN32
     // TODO: VirtualAlloc(..., MEM_RESET, ...) may be faster
     // See https://arvid.io/2018/04/02/memory-mapping-on-windows/#1
     // Not sure what the differences are
-    if (!VirtualFree(getFrame(frameIdx), frameSize, MEM_DECOMMIT)) {
+    if (!VirtualFree(discardGranule, discardGranuleSize, MEM_DECOMMIT)) {
         auto code = GetLastError();
         throw BufferManagerException(std::format(
             "Releasing physical memory associated with a frame failed with error code {}: {}.",
             code, systemErrMessage(code)));
     }
 
 #else
-    int error = madvise(getFrame(frameIdx), frameSize, MADV_DONTNEED);
+    int error = madvise(discardGranule, discardGranuleSize, MADV_DONTNEED);
     if (error != 0) {
         // LCOV_EXCL_START
         throw BufferManagerException(std::format(
@@ -80,6 +119,7 @@ void VMRegion::releaseFrame(frame_idx_t frameIdx) const {
         // LCOV_EXCL_STOP
     }
 #endif
+    return discardGranuleSize;
 }
 
 frame_group_idx_t VMRegion::addNewFrameGroup() {
@@ -89,8 +129,23 @@ frame_group_idx_t VMRegion::addNewFrameGroup() {
         throw BufferManagerException("No more frame groups can be added to the allocator.");
         // LCOV_EXCL_STOP
     }
+    auto residentFrameCounts =
+        std::make_unique<std::atomic<uint16_t>[]>(numDiscardGranulesPerFrameGroup);
+    for (auto i = 0u; i < numDiscardGranulesPerFrameGroup; ++i) {
+        residentFrameCounts[i].store(0);
+    }
+    residentFramesPerDiscardGranule.push_back(std::move(residentFrameCounts));
     return numFrameGroups++;
 }
 
+uint64_t VMRegion::getFrameGroupIdx(frame_idx_t frameIdx) const {
+    return frameIdx >> StorageConstants::PAGE_GROUP_SIZE_LOG2;
+}
+
+uint64_t VMRegion::getDiscardGranuleIdxInFrameGroup(frame_idx_t frameIdx) const {
+    const auto frameIdxInGroup = frameIdx & StorageConstants::PAGE_IDX_IN_GROUP_MASK;
+    return frameIdxInGroup * static_cast<uint64_t>(frameSize) / discardGranuleSize;
+}
+
 } // namespace storage
 } // namespace lbug