自从在 CMU 的 DB Talk 里看到了 DuckDB，我就对这个数据库很感兴趣。加上最近发出来的两篇很有意思的博客（西方最快的排序、DuckDB in Browser），让我按捺不住源码阅读的心，于是准备自底向上看一下它的一些实现细节，以及思考一下这些点对 TiDB 的开发是否有帮助。
DuckDB 是一个单机分析数据库，他们宣传的 SQLite for Analytics 已经说明了一切。很多数据分析师所需要分析的数据通常单机就能存下，然而一方面她们需要使用 python 写很多数据清洗和预处理的逻辑，堪称重新发明 SQL；另一方面磁盘能存下不代表内存能放下，如果要一次性分析超过内存空间大小的数据集，分析师们则不得不手动切分数据。在这个基础上，用支持 SQL 和内存管理的数据库来辅助，就显得尤为香甜。
服务于这个目标，DuckDB 一方面实现了一个高性能的列存计算引擎，另一方面采用了精巧的内存管理机制来避免 OOM。前者可待未来的讲解，这篇文章我们着重看后者。
DuckDB 的缓存管理器实现了两个重要功能。第一个功能就是传统数据库中数据页缓存管理，这与前文提到的一致；第二个功能则是对运行时产生的中间数据进行管理，譬如排序或者 Join 的结果。这些结果不需要保存下来，因此原始而直接地放在内存中是很多数据库采取的管理手段。然而当数据库中有大量并发查询占用内存，则数据库整体的内存占用将失去控制，导致被操作系统的 OOM Killer 杀掉。
这个问题对于 TiDB 来说，已经给用户与研发造成了习得性无助。因此这里我们着重关注第二个功能的实现。
毫无意外地，DuckDB 的缓存管理器是由一个名为 BufferManager 的类实现的。这个类最重要的几个公共方法是
//! Register a block with the given block id in the base file
shared_ptr<BlockHandle> RegisterBlock(block_id_t block_id);
//! Register an in-memory buffer of arbitrary size, as long as it is >= BLOCK_SIZE. can_destroy signifies whether or
//! not the buffer can be destroyed when unpinned, or whether or not it needs to be written to a temporary file so
//! it can be reloaded. The resulting buffer will already be allocated, but needs to be pinned in order to be used.
shared_ptr<BlockHandle> RegisterMemory(idx_t block_size, bool can_destroy);
//! Allocate an in-memory buffer with a single pin.
//! The allocated memory is released when the buffer handle is destroyed.
unique_ptr<BufferHandle> Allocate(idx_t block_size);
unique_ptr<BufferHandle> Pin(shared_ptr<BlockHandle> &handle);BufferManager 通过 RegisterBlock 和 RegisterMemory 返回 shared_ptr<BlockHandle>，前者申请的 BlockHandle 实现了类似传统数据库中的缓存管理器的职能，而后者是对运行时内存的管理。
DuckDB 中的 Block 就是缓存管理器管理的缓存单元，类似页的概念。有两种 Block，一种是 block id 小于 MAXIMUM_BLOCK （其值等于 1 << 62）的 Block，代表磁盘中的页；另一种是 block id 大于MAXIMUM_BLOCK 的 Block，则代表是运行时产生的临时 Block。RegisterBlock 方法并不会分配新的 block id，而是传入一个已经在磁盘上的 block 的 block id，表示请缓存管理器来管理这个 block。而 RegisterMemory 不需要 block id，而是接收两个参数，一个 block_size 用于指定分配的 block 的大小，另一个 can_destroy 表示这个 block 是否可以在暂时不用时直接释放，主要是在 Allocate 这个方法调用时被设置为 true，即一次性使用的 block。Pin 方法就是之前介绍过的，表示把 Block 固定在内存中。
shared_ptr<BlockHandle> BufferManager::RegisterMemory(idx_t block_size, bool can_destroy) {
auto alloc_size = block_size + Storage::BLOCK_HEADER_SIZE;
// first evict blocks until we have enough memory to store this buffer
if (!EvictBlocks(alloc_size, maximum_memory)) {
throw OutOfMemoryException("could not allocate block of %lld bytes%s", alloc_size, InMemoryWarning());
}
// allocate the buffer
auto temp_id = ++temporary_id;
auto buffer = make_unique<ManagedBuffer>(db, block_size, can_destroy, temp_id);
// create a new block pointer for this block
return make_shared<BlockHandle>(db, temp_id, move(buffer), can_destroy, block_size);
}具体实现的逻辑并不复杂。首先是检查内存限制，如果新增这块内存会导致内存超限，则首先驱逐掉暂时不用的一些页。然后申请缓存空间，这里的原子变量 temporary_id 初始化的值就是 MAXIMUM_BLOCK，这里自增分配可以保证通过 RegisterMemory 注册的 block id 都大于 MAXIMUM_BLOCK，从而区分数据库本身的 block 与运行时的 block。ManagedBuffer 继承自 FileBuffer，FileBuffer 是数据真实存放的地方，在初始化时会在内存中开辟所需的 block_size 的空间，以便执行器直接使用，而在被落盘后，提供了面向磁盘进行读写的接口。最后，将这个 ManagedBuffer 移入 BlockHandle 中，交由 BlockHandle 进行管理。
BlockHandle 有两种状态，已经加载和尚未加载。RegisterBlock 并不会主动加载 Block，而 RegisterMemory 则是直接在内存中开辟的空间，因此返回的 Block 是已经加载的状态。
// RegisterBlock 调用的 BlockHandle 的 constructor
BlockHandle::BlockHandle(DatabaseInstance &db, block_id_t block_id_p)
: db(db), readers(0), block_id(block_id_p), buffer(nullptr), eviction_timestamp(0), can_destroy(false) {
eviction_timestamp = 0;
state = BlockState::BLOCK_UNLOADED;
memory_usage = Storage::BLOCK_ALLOC_SIZE;
}
// RegisterMemory 调用的 BlockHandle 的 constructor
BlockHandle::BlockHandle(DatabaseInstance &db, block_id_t block_id_p, unique_ptr<FileBuffer> buffer_p,
bool can_destroy_p, idx_t block_size)
: db(db), readers(0), block_id(block_id_p), eviction_timestamp(0), can_destroy(can_destroy_p) {
D_ASSERT(block_size >= Storage::BLOCK_SIZE);
buffer = move(buffer_p);
state = BlockState::BLOCK_LOADED;
memory_usage = block_size + Storage::BLOCK_HEADER_SIZE;
}因此 BlockHandle 提供了加载与卸载的两个方法用于进行状态转换。
unique_ptr<BufferHandle> BlockHandle::Load(shared_ptr<BlockHandle> &handle) {
if (handle->state == BlockState::BLOCK_LOADED) {
// already loaded
D_ASSERT(handle->buffer);
return make_unique<BufferHandle>(handle, handle->buffer.get());
}
auto &buffer_manager = BufferManager::GetBufferManager(handle->db);
auto &block_manager = BlockManager::GetBlockManager(handle->db);
if (handle->block_id < MAXIMUM_BLOCK) {
auto block = make_unique<Block>(Allocator::Get(handle->db), handle->block_id);
block_manager.Read(*block);
handle->buffer = move(block);
} else {
if (handle->can_destroy) {
return nullptr;
} else {
handle->buffer = buffer_manager.ReadTemporaryBuffer(handle->block_id);
}
}
handle->state = BlockState::BLOCK_LOADED;
return make_unique<BufferHandle>(handle, handle->buffer.get());
}首先看 Load 的实现。当 Block 的持有者需要使用这块 Block 时，可以直接调用 Load 获得一个 BufferHandle，这个 BufferHandle 是对 FileBuffer 做的简单封装，其生命周期从构造时调用 Pin 方法开始到析构时自动 Unpin 结束，从而使用者可以无需操心 Pin/Unpin 的流程。如果当前 Block 是尚未加载的状态，则通过 block id 判断是两类 block 分别处理。
对于 handle->block_id < MAXIMUM_BLOCK 的数据页，通过创建一个 Block 类申请一块内存空间，这里的 Block 也继承自 FileBuffer，是根正苗红的 Block。然后由 block_manager 负载加载这块 Block。BlockManager 有两类实现，一个是 SIngleFileBlockManager，另一个是 InMemoryBlockManager，分别负责在两种 DuckDB 运行状态下的数据管理，前者是数据会持久化到磁盘中的一个文件（即用 .open 打开的那个文件，类似 SQLite），而后者则是纯内存运行，此处不细说。
对于另一种运行时 Block，如果这个 Block 是可以被销毁的，则只能在创建时被 Load 一次，释放过后再次 Load 就无效了，因此直接返回 nullptr，否则调用 BufferManager 提供的 ReadTemporaryBuffer 方法来加载之前 Unload 时保存下来的数据。
void BlockHandle::Unload() {
auto &buffer_manager = BufferManager::GetBufferManager(db);
if (state == BlockState::BLOCK_UNLOADED) {
// already unloaded: nothing to do
return;
}
D_ASSERT(CanUnload());
D_ASSERT(memory_usage >= Storage::BLOCK_ALLOC_SIZE);
if (block_id >= MAXIMUM_BLOCK && !can_destroy) {
// temporary block that cannot be destroyed: write to temporary file
buffer_manager.WriteTemporaryBuffer((ManagedBuffer &)*buffer);
}
buffer.reset();
buffer_manager.current_memory -= memory_usage;
state = BlockState::BLOCK_UNLOADED;
}Unload 的实现更简单。如果已经 Unload 了，自然不用再操作，否则这里只需要对不能直接销毁的运行时页进行保存，调用 WriteTempraryBuffer 方法即可。注意到这里对内存占用的修改，原子变量 current_memory 表示整个 BufferManager 管理的数据所占内存空间的大小。
Read/WriteTemporaryBuffer 的实现也非常简单，直接用 block id 作为文件名，把这个 block 写入到临时文件目录中即可，只是记得在 Read 完成后删除文件就好。
以上就是 DuckDB 缓存管理器实现的分析，可以看到，通过区分两类 Block，DuckDB 可以把磁盘上的数据页和运行时需要申请的空间都管理起来，并限制住内存占用上限，并在达到上限时自动地落盘释放掉一些 Block。
我们可以继续看 RegisterMemory 和 Allocate 这两个申请运行时内存的接口的使用的例子。
DuckDB 作为分析型数据库，自然首选列的形式来表示数据。一个关键的点在于，使用列的形式是为了更好地利用 cache 以及 SIMD 。但除了表达式计算外，对于聚合、hash join 和排序这类需要改变行的数量或者顺序的操作，其实按行来处理会更好。譬如对于排序来说，比较大小以行为基础的，而 reorder 在逻辑上是以也行为单位的，因此 DuckDB 还引入了行形式的内存表示。
根据数据类型，数据有固定大小和大小可变的两类。如整形、浮点数等类型就是固定大小的，而字符串、二进制等数据类型就是大小可变的。DuckDB 在进行列转行时，会把固定大小的列保留，而大小可变的列则堆积放入另外一块内存空间（StringHeap），在原来的位置上保存指向 StringHeap 中的指针就好了。这样转换后，一行的数据宽度就是固定的了。
行数据以 RowDataBlock 类形式保存，一个 RowDataBlock 可以包含每行 entry_size 宽度的 capacity 数量的行，因此构造函数直接调用缓存管理器的 RegisterMemory 方法申请 capacity * entry_size 大小的 block 即可。返回的 BlockHandler 以成员变量的形式保存在类中。
struct RowDataBlock {
RowDataBlock(BufferManager &buffer_manager, idx_t capacity, idx_t entry_size)
: capacity(capacity), entry_size(entry_size), count(0), byte_offset(0) {
block = buffer_manager.RegisterMemory(capacity * entry_size, false);
}
//! The buffer block handle
shared_ptr<BlockHandle> block;
//! Capacity (number of entries) and entry size that fit in this block
idx_t capacity;
const idx_t entry_size;
//! Number of entries currently in this block
idx_t count;
//! Write offset (if variable size entries)
idx_t byte_offset;
};当执行器需要访问这块空间时，调用缓存管理器的 Pin 方法即可获得一个 BufferHandler，其中就包含这块空间在内存中的地址，从而可以随意地进行读写。而当执行器暂时不使用这块空间时，释放掉 BufferHandler 即可。而缓存管理器则会在背后默默地控制内存的占用。
TiDB 目前内存管控的机制是在各个算子内部统计申请的 chunk 的内存占用，以 session 为单位记录总共的内存占用，直到统计的内存超出限制后，使用名为 RowContainer 的结构实现部分 chunk 的落盘。
在我看来，这套机制的弊端非常明显。
如果能在 TiDB 中引入 session 级别的缓存管理器，从 chunk、column 等数据结构的分配源头就做内存追踪与统计，并在内存使用量达到限度后自动且无感知地 spill，将会是非常美好的未来。
具体的细节，后续可能会有设计文档，欢迎大家留言讨论~