Linux kernel rbtree

Posted 2020-10-21 舰队技术博客

Linux kernel rbtree

因编写内核模块时需要用到rbtree来记录异步request，研究分析了一下kernel rbtree的使用方法，记录于此。本文主要参考了内核文档rbtree.txt

rbtree简介

Red-black trees（rbtree）是一种自平衡的二叉搜索树，用于存储可分类的key/value数据对。它不同于radix trees或者hash tables。
radix trees用于有效存储稀疏数组（使用长整型索引进行节点的插入、查询和删除），其索引值太大无法用数组直接存储。
hash tables用于散列索引缩小查询的范围，但它没有做排序，因此不能快速的定位。

Red-black trees和AVL trees很相似，但是提供了最坏情况下更快的实时插入和删除性能。插入最多2次rotations、删除最多3次rotations即可完成tree的重平衡。不过相比AVL trees，其查询时间稍慢（O(log n)）。

Linux内核大量使用rbtree，如：I/O调度算法deadline和CFQ使用rbtree来跟踪request；高精度定时器代码使用rbtree来组织定时任务；ext3文件系统使用rbtree来跟踪目录entry；等等。

rbtree使用方法

内核rbtree的实现在文件"lib/rbtree.c"，使用rbtree需要包含头文件：

#include <linux/rbtree.h>

为了提高性能，linux rbtree比传统的tree实现了更少的中间层。rbtree的节点结构体struct rb_node直接嵌入到使用者的data structure（传统的方法是通过指针指向了data structure）。rbtree的插入和查询函数由使用者通过调用linux rbtree提供的基础函数自己实现（传统的方法是提供回调函数指针）。并且btree的锁也由使用者自己管理。

创建rbtree

在data数据结构里定义struct rb_node：

struct mytype {
    struct rb_node node;
    char *keystring;
};

当处理rbtree的节点时，通过container_of()宏定义找到data数据结构指针。keystring为rbtree的key，可以定义为字符串或者整型，它将用于用户自定义的排序和查找。

然后定义rbtree的root节点：

struct rb_root mytree = RB_ROOT;

查找rbtree

使用者自己实现rbtree的查找函数，通过如下方法：从root开始，比较key的值，然后根据需要查找left节点或者right节点。

struct mytype *my_search(struct rb_root *root, char *string)
{
    struct rb_node *node = root->rb_node;
    while (node) {
        struct mytype *data = container_of(node, struct mytype, node);
        int result;
        result = strcmp(string, data->keystring);
        if (result < 0)
            node = node->rb_left;
        else if (result > 0)
            node = node->rb_right;
        else
            return data;
    }
    return NULL;
}

插入新节点

使用者自己实现rbtree的插入函数，先找到插入的位置（该位置为NULL），然后插入新的节点并执行rbtree的重平衡。在查找到插入位置时，需要其parent节点的link用于rbtree的重平衡。

int my_insert(struct rb_root *root, struct mytype *data)
{
    struct rb_node **new = &(root->rb_node), *parent = NULL;
    /* Figure out where to put new node */
    while (*new) {
        struct mytype *this = container_of(*new, struct mytype, node);
        int result = strcmp(data->keystring, this->keystring);
        parent = *new;
        if (result < 0)
            new = &((*new)->rb_left);
        else if (result > 0)
            new = &((*new)->rb_right);
        else
            return FALSE;
    }
    /* Add new node and rebalance tree. */
    rb_link_node(&data->node, parent, new);
    rb_insert_color(&data->node, root);
    return TRUE;
}

删除/覆盖节点

通过如下函数删除和覆盖一个节点：

void rb_erase(struct rb_node *victim, struct rb_root *tree);
void rb_replace_node(struct rb_node *old, struct rb_node *new, struct rb_root *tree);

覆盖一个节点并不会重平衡rbtree，因此必须保证new和old的key是一样的，否者会导致异常。
删除一个节点代码示例：

struct mytype *data = mysearch(&mytree, "walrus");
if (data) {
    rb_erase(&data->node, &mytree);
    myfree(data);
}

按顺序遍历rbtree

如下4个函数用于顺序遍历rbtree：

struct rb_node *rb_first(struct rb_root *tree);
struct rb_node *rb_last(struct rb_root *tree);
struct rb_node *rb_next(struct rb_node *node);
struct rb_node *rb_prev(struct rb_node *node);

代码示例：

struct rb_node *node;
for (node = rb_first(&mytree); node; node = rb_next(node))
    printk("key=%s\n", rb_entry(node, struct mytype, node)->keystring);