CMU15445 (Fall 2020) 数据库系统 Project#2 - B+ Tree 详解(下篇)
阅读原文时间:2023年08月09日阅读:1

前言

上一篇博客中实现了单线程 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 实现对应的方法,似乎也没那么难以下手了。完成之后可以明显感受到精神力得以增强,信心开始膨胀(不是。

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

手机扫一扫

移动阅读更方便

阿里云服务器
腾讯云服务器
七牛云服务器