CMU15445 (Fall 2020) 数据库系统 Project#2 - B+ Tree 详解（下篇）

前言

上一篇博客中实现了单线程 B+ 树的查找、插入、删除和迭代操作，这篇博客将完成实验二的剩余任务：并发 B+ 树。实现 B+ 树并发访问最简单的方法就是在整棵树上加一把大锁，但是这样会导致过多线程处于阻塞状态，严重降低 B+ 树的性能。这篇博客将使用蟹行协议（crabbing protocol）实现并发。

蟹行协议

该协议的名字来源于螃蟹走路的方式：先移动一边的腿，然后另一边的，如此交替进行。该协议的加锁过程，从上往下和从下往上（发生分裂、合并或重新分布的情况），就像螃蟹移动一样。

查找

当查找一个键时，蟹行协议首先用共享模式锁住根结点。沿树向下遍历，在子结点上获得锁以后，它释放父结点上的锁。它重复该过程直至叶结点。

实际上在对根节点上锁之前，还需要对根节点的 id 进行上锁，防止根节点发生变化。所以需要在 BPlusTree 中添加一个 std::mutex root_latch_ 成员，在查找、插入、删除和迭代之前都需要获取 root_latch_。

插入和删除

插入和删除都需要对节点加写锁，由于插入可能导致叶节点分裂，删除可能导致叶节点的合并或者重新分配，所以在释放父节点的锁之前需要判断子节点是否安全。只有安全时才能释放所有祖先节点的锁，否则需要一直加锁下去。

插入时只要节点的键值对数量小于 max_internal_size_ - 1 （最后一个键值对充当哨兵），就不会分裂，就是安全的。

删除时需要特殊处理根节点：

• 根节点为叶节点，最小的键值对数量是 1，所以删除之前需要键值对数量大于 1 才是安全的；
• 根节点为叶节点，最小的键值对数量是 2，所以删除之前需要键值对数量大于 2 才是安全的；

如果节点不是根节点，需要删除后仍处于半满状态才是安全的。

INDEX_TEMPLATE_ARGUMENTS
bool BPLUSTREE_TYPE::IsPageSafe(BPlusTreePage *page, OperationType operation) {
  auto size = page->GetSize();

  switch (operation) {
    case OperationType::READ:
      return true;

    case OperationType::INSERT:
      return size < page->GetMaxSize() - 1;

    case OperationType::REMOVE:
      if (page->IsRootPage()) {
        return page->IsLeafPage() ? size > 1 : size > 2;
      }

      return size > page->GetMinSize();

    default:
      break;
  }

  return false;
}

以删除为例，由于下图的 B 节点删除后不满足半满状态，所以不安全，无法释放 A 上的锁。

当走到 D 节点时，发现 D 是安全的，这时候可以释放所有祖先节点（A 和 B）上的锁。

对子节点的加锁发生在 FindLeafPage() 函数中，当子节点不安全时，调用 Transaction::AddIntoPageSet(Page *) 记录父节点。对于 root_latch_，当根节点不安全时，加到 transaction 里面的是空指针：

INDEX_TEMPLATE_ARGUMENTS
Page *BPLUSTREE_TYPE::FindLeafPage(const KeyType &key, bool leftMost, OperationType operation,
                                   Transaction *transaction) {
  if (operation == OperationType::READ) {
    root_latch_.lock();
  }

  auto page_id = root_page_id_;
  auto page = buffer_pool_manager_->FetchPage(page_id);
  auto node = ToTreePage(page);

  // 给根节点上锁
  if (operation == OperationType::READ) {
    page->RLatch();
    root_latch_.unlock();
  } else {
    page->WLatch();
    if (!IsPageSafe(node, operation)) {
      transaction->AddIntoPageSet(nullptr);  // 加一个空指针表示根节点 id 的锁
    } else {
      root_latch_.unlock();
    }
  }

  // 定位到包含 key 的叶节点
  while (!node->IsLeafPage()) {
    InternalPage *inode = ToInternalPage(node);

    // 寻找下一个包含 key 的节点
    if (!leftMost) {
      page_id = inode->Lookup(key, comparator_);
    } else {
      page_id = inode->ValueAt(0);
    }

    // 移动到子节点
    auto child_page = buffer_pool_manager_->FetchPage(page_id);

    // 给子节点上锁
    if (operation == OperationType::READ) {
      child_page->RLatch();
      page->RUnlatch();
      buffer_pool_manager_->UnpinPage(page->GetPageId(), false);
    } else {
      child_page->WLatch();
      transaction->AddIntoPageSet(page);

      // 如果子节点安全，就释放所有祖先节点上的写锁
      if (IsPageSafe(ToTreePage(child_page), operation)) {
        UnlockAncestors(transaction);
      }
    }

    page = child_page;
    node = ToTreePage(page);
  }

  return page;
}

如果子节点按钮，调用 UnlockAncestors() 来释放祖先节点上的锁，注意这里是先解锁再 Unpin()，如果先 Unpin() 可能导致在解锁之前页被换出，这时候解锁的是别人的页了：

INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::UnlockAncestors(Transaction *transaction, bool unpin) {
  auto pages = transaction->GetPageSet().get();

  while (!pages->empty()) {
    auto page = pages->front();
    pages->pop_front();

    if (!page) {
      root_latch_.unlock();
    } else {
      page->WUnlatch();
      if (unpin) {
        buffer_pool_manager_->UnpinPage(page->GetPageId(), false);
      }
    }
  }
}

查找

GetValue() 函数修改如下，在判断树是否为空之前需要加锁：

INDEX_TEMPLATE_ARGUMENTS
bool BPLUSTREE_TYPE::GetValue(const KeyType &key, std::vector<ValueType> *result, Transaction *transaction) {
  root_latch_.lock();
  if (IsEmpty()) {
    root_latch_.unlock();
    return false;
  }
  root_latch_.unlock();

  // 在叶节点中寻找 key
  auto leaf_page = FindLeafPage(key);
  LeafPage *leaf = ToLeafPage(leaf_page);

  ValueType value;
  auto success = leaf->Lookup(key, &value, comparator_);
  if (success) {
    result->push_back(value);
  }

  leaf_page->RUnlatch();
  buffer_pool_manager_->UnpinPage(leaf->GetPageId(), false);
  return success;
}

插入

在插入之前需要对 root_latch_ 上锁，在结束之前需要释放所有祖先节点上的锁：

INDEX_TEMPLATE_ARGUMENTS
bool BPLUSTREE_TYPE::Insert(const KeyType &key, const ValueType &value, Transaction *transaction) {
  root_latch_.lock();
  // 省略部分代码
}

/* Insert constant key & value pair into an empty tree */
INDEX_TEMPLATE_ARGUMENTS void BPLUSTREE_TYPE::StartNewTree(const KeyType &key, const ValueType &value) {
  // 创建一个叶节点作为根节点，并插入新数据
  // 省略部分代码

  UpdateRootPageId(1);
  root_latch_.unlock();
  buffer_pool_manager_->UnpinPage(root_page_id_, true);
}

/* Insert constant key & value pair into leaf page */
INDEX_TEMPLATE_ARGUMENTS
bool BPLUSTREE_TYPE::InsertIntoLeaf(const KeyType &key, const ValueType &value, Transaction *transaction) {
  // 定位到包含 key 的叶节点
  auto leaf_page = FindLeafPage(key, false, OperationType::INSERT, transaction);
  LeafPage *leaf = ToLeafPage(leaf_page);

  // 不能插入相同的键
  ValueType exist_value;
  if (leaf->Lookup(key, &exist_value, comparator_)) {
    UnlockAncestors(transaction);
    leaf_page->WUnlatch();
    buffer_pool_manager_->UnpinPage(leaf->GetPageId(), false);
    return false;
  }

  // 省略部分代码
  UnlockAncestors(transaction);
  leaf_page->WUnlatch();
  buffer_pool_manager_->UnpinPage(leaf->GetPageId(), true);
  return true;
}

/* Insert key & value pair into internal page after split */
INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::InsertIntoParent(BPlusTreePage *old_node, const KeyType &key, BPlusTreePage *new_node,
                                      Transaction *transaction) {
  // 根节点发生分裂需要新建一个根节点，B+树的高度 +1
  if (old_node->IsRootPage()) {
    // 省略部分代码
    UpdateRootPageId(0);
    buffer_pool_manager_->UnpinPage(root_page_id_, true);
    UnlockAncestors(transaction, false);
    return;
  }

  // 省略部分代码
  // 父节点溢出时需要再次分裂
  if (size == internal_max_size_) {
    // 省略
  } else {
    UnlockAncestors(transaction, false);
    buffer_pool_manager_->UnpinPage(parent_id, true);
  }
}

删除

删除和插入类似，唯一需要注意的是对兄弟节点进行加锁，防止迭代的时候被访问，调整结束后立即释放兄弟节点上的锁：

INDEX_TEMPLATE_ARGUMENTS
void BPLUSTREE_TYPE::Remove(const KeyType &key, Transaction *transaction) {
  root_latch_.lock();

  if (IsEmpty()) {
    root_latch_.unlock();
    return;
  }

  // 定位到叶节点并删除键值对
  auto leaf_page = FindLeafPage(key, false, OperationType::REMOVE, transaction);
  LeafPage *leaf = ToLeafPage(leaf_page);
  auto old_size = leaf->GetSize();
  auto size = leaf->RemoveAndDeleteRecord(key, comparator_);

  // 叶节点删除之后没有处于半满状态需要合并相邻节点或者重新分配
  if (size < leaf->GetMinSize() && CoalesceOrRedistribute(leaf, transaction)) {
    transaction->AddIntoDeletedPageSet(leaf->GetPageId());
  }

  UnlockAncestors(transaction);
  leaf_page->WUnlatch();
  buffer_pool_manager_->UnpinPage(leaf->GetPageId(), old_size != size);

  // 不知道为什么删除之后会导致堆溢出错误
  // DeletePages(transaction);
}

INDEX_TEMPLATE_ARGUMENTS
template <typename N>
bool BPLUSTREE_TYPE::CoalesceOrRedistribute(N *node, Transaction *transaction) {
  // 找到相邻的兄弟节点并加锁，省略部分代码
  auto sibling_page = buffer_pool_manager_->FetchPage(parent->ValueAt(sibling_index));
  N *sibling = reinterpret_cast<N *>(sibling_page->GetData());
  sibling_page->WLatch();

  // 如果两个节点的大小和大于 max_size-1，就直接重新分配，否则直接合并兄弟节点
  bool is_merge = sibling->GetSize() + node->GetSize() <= node->GetMaxSize() - 1;
  if (is_merge) {
    Coalesce(&sibling, &node, &parent, index, transaction);
  } else {
    Redistribute(sibling, node, index);
  }

  // 兄弟节点解锁
  sibling_page->WUnlatch();
  buffer_pool_manager_->UnpinPage(parent->GetPageId(), true);
  buffer_pool_manager_->UnpinPage(sibling->GetPageId(), true);
  return is_merge;
}

迭代

迭代时需要对叶节点加读锁，析构迭代器时需要释放锁：

INDEX_TEMPLATE_ARGUMENTS
INDEXITERATOR_TYPE &INDEXITERATOR_TYPE::operator++() {
  if (isEnd()) {
    return *this;
  }

  LeafPage *leaf = reinterpret_cast<LeafPage *>(page_);
  if (index_ < leaf->GetSize() - 1) {
    index_++;
  } else {
    Page* old_page = page_;

    // 移动到下一页
    page_id_ = leaf->GetNextPageId();
    if (page_id_ != INVALID_PAGE_ID) {
      page_  = buffer_pool_manager_->FetchPage(page_id_);
      page_->RLatch();
    } else {
      page_ = nullptr;
    }

    index_ = 0;
    old_page->RUnlatch();
    buffer_pool_manager_->UnpinPage(old_page->GetPageId(), false);
  }

  return *this;
}

INDEX_TEMPLATE_ARGUMENTS
INDEXITERATOR_TYPE::~IndexIterator() {
  if (!isEnd()) {
    page_->RUnlatch();
    buffer_pool_manager_->UnpinPage(page_->GetPageId(), false);
  }
};

测试

在终端输入下述指令完成编译：

cd build
cmake ..
make

# 从 Grade scope 拔下来的测试用例
make b_plus_tree_checkpoint_2_concurrent_test
make b_plus_tree_bench_test

./test/b_plus_tree_checkpoint_2_concurrent_test
./test/b_plus_tree_bench_test

测试结果如下，只给虚拟机分配了一个核，速度可能慢了一些：

后记

纸上得来终觉浅，绝知此事要躬行。历经三天终于完成了 B+ 树的实验，通过这次实验可以加深对索引结构的理解，可以说是非常硬核的一次实验了。刚开始有点无从下手，因为要完成的函数实在太多了。写着写着发现可以自顶而下，先完成 BPlusTree 方法上的逻辑，再深入到底层的 BPlusTreePage 实现对应的方法，似乎也没那么难以下手了。完成之后可以明显感受到精神力得以增强，信心开始膨胀（不是。

借用屑老板的话：这是一场「试炼」，我认为这就是一场为了战胜过去的「试炼」，只有在战胜了那些幼稚的过去，人才能有所成长。嗯？你也是那样的吧？