in_mem_overflow_buffer.h

Go to the documentation of this file.

#pragma once
 
#include <iterator>
#include <memory>
#include <vector>
 
#include "api.h"
#include "copy_constructors.h"
 
namespace kuzu {
namespace storage {
class MemoryBuffer;
class MemoryManager;
} // namespace storage
 
namespace common {
 
struct KUZU_API BufferBlock {
public:
    explicit BufferBlock(std::unique_ptr<storage::MemoryBuffer> block);
    ~BufferBlock();
 
    uint64_t size() const;
    uint8_t* data() const;
 
public:
    uint64_t currentOffset;
    std::unique_ptr<storage::MemoryBuffer> block;
 
    void resetCurrentOffset() { currentOffset = 0; }
};
 
class KUZU_API InMemOverflowBuffer {
 
public:
    explicit InMemOverflowBuffer(storage::MemoryManager* memoryManager)
        : memoryManager{memoryManager} {};
 
    DEFAULT_BOTH_MOVE(InMemOverflowBuffer);
 
    uint8_t* allocateSpace(uint64_t size);
 
    void merge(InMemOverflowBuffer& other) {
        move(begin(other.blocks), end(other.blocks), back_inserter(blocks));
        // We clear the other InMemOverflowBuffer's block because when it is deconstructed,
        // InMemOverflowBuffer's deconstructed tries to free these pages by calling
        // memoryManager->freeBlock, but it should not because this InMemOverflowBuffer still
        // needs them.
        other.blocks.clear();
    }
 
    // Releases all memory accumulated for string overflows so far and re-initializes its state to
    // an empty buffer. If there is a large string that used point to any of these overflow buffers
    // they will error.
    void resetBuffer();
 
    // Manually set the underlying memory buffer to evicted to avoid double free
    void preventDestruction();
 
    storage::MemoryManager* getMemoryManager() { return memoryManager; }
 
private:
    bool requireNewBlock(uint64_t sizeToAllocate) {
        return blocks.empty() ||
               (currentBlock()->currentOffset + sizeToAllocate) > currentBlock()->size();
    }
 
    void allocateNewBlock(uint64_t size);
 
    BufferBlock* currentBlock() { return blocks.back().get(); }
 
private:
    std::vector<std::unique_ptr<BufferBlock>> blocks;
    storage::MemoryManager* memoryManager;
};
 
} // namespace common
} // namespace kuzu