内核红黑树-I

内核中许多部分使用到红黑树,比如tcp的乱序队列out_of_order_queue,TC中fq公平队列的流表等。以下为红黑树的5个准则:

  1. A node is either red or black
    节点的颜色非红即黑
  2. The root is black
    根节点为黑色
  3. All leaves (NULL) are black
    所有的叶子(NULL)节点都为黑色
  4. Both children of every red node are black
    红色节点的子节点都为黑色
  5. Every simple path from root to leaves contains the same number of black nodes.
    每个由根到叶子节点的路径必须包含相同数量的黑色节点

以下据此看看内核的红黑树实现。

1 红黑树遍历

以下以fq公平队列为例,看一下红黑树遍历。

在函数fq_classify中,需要遍历红黑树查找对应的流结构,这里使用函数rb_entry函数,其等价于container_of函数由node找到对应的fq_flow结构。红黑树的键值为套接口结构sk指针的数值,小于等于此sk值的流结构在左树,大于此值的流结构位于右树。

static struct fq_flow *fq_classify(struct sk_buff *skb, struct fq_sched_data *q)
{
    struct rb_node **p, *parent;
    struct sock *sk = skb->sk;
    struct rb_root *root;
    struct fq_flow *f;
    ...
    p = &root->rb_node;
    parent = NULL;
    while (*p) {
        parent = *p;

        f = rb_entry(parent, struct fq_flow, fq_node);
        if (f->sk == sk) {
            ...
        }
        if (f->sk > sk)
            p = &parent->rb_right;
        else
            p = &parent->rb_left;
    }

2 节点插入

节点的插入分为两个部分:节点插入和红黑色颜色的平衡

同样是函数fq_classify,调用函数rb_link_node执行节点插入,其中参数p表示的为父节点的左子节点或者右子节点。

static struct fq_flow *fq_classify(struct sk_buff *skb, struct fq_sched_data *q)
{
    struct rb_node **p, *parent;
    struct sock *sk = skb->sk;
    struct rb_root *root;
    struct fq_flow *f;
    ...
    f = kmem_cache_zalloc(fq_flow_cachep, GFP_ATOMIC | __GFP_NOWARN);
    ...
    f->sk = sk;
    ...

    rb_link_node(&f->fq_node, parent, p);
    rb_insert_color(&f->fq_node, root);

2.1 节点插入操作

函数rb_link_node用于节点的插入,即将node节点插入到parent节点的左树或者右树(由rb_link表示),另外,在节点的__rb_parent_color字段设置父节点指针,并且将节点的左右节点都设置为空NULL。

static inline void rb_link_node(struct rb_node *node, struct rb_node *parent,
                struct rb_node **rb_link)
{       
    node->__rb_parent_color = (unsigned long)parent;
    node->rb_left = node->rb_right = NULL;
    
    *rb_link = node;
} 

2.2 红黑树平衡

最重要的是以下的rb_insert_color函数,其封装了函数__rb_insert,用于处理以上的节点插入操作之后,红黑树的颜色调整等。

void rb_insert_color(struct rb_node *node, struct rb_root *root)
{       
    __rb_insert(node, root, false, NULL, dummy_rotate);
} 

以下为__rb_insert函数,while循环结束的条件为parent节点为黑色,如果其为红色将一直循环。

static __always_inline void
__rb_insert(struct rb_node *node, struct rb_root *root,
        bool newleft, struct rb_node **leftmost,
        void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
{
    struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;

    if (newleft)
        *leftmost = node;

    while (true) {
        /*
         * Loop invariant: node is red.
         */

2.2.1 插入根节点情况

如果红黑树为空,根据准则2,新插入节点为根节点,颜色为黑色。另外还有一种情况,以下将会看到,在Case-1的时候,需要由红黑树的底部往上检查是否符合准则4,parent为空时表示检查完毕。

        if (unlikely(!parent)) {
            /*
             * The inserted node is root. Either this is the first node, 
             * or we recursed at Case 1 below and are no longer violating 4).
             */
            rb_set_parent_color(node, NULL, RB_BLACK);
            break;
        }

2.2.2 父节点为黑色情况

如果父节点parent的颜色为黑色,同时由于新插入节点的颜色为红色,并不会违反红黑树准则(特别是第5准则),插入完成。但是如果父节点为红色,需要注意准则4,红色节点的子节点为黑色。

        /*
         * If there is a black parent, we are done.
         * Otherwise, take some corrective action as,
         * per 4), we don't want a red root or two
         * consecutive red nodes.
         */
        if(rb_is_black(parent))
            break;

2.2.3 父节点为红色情况-I

以下处理父节点parent为红色的情况,根据红黑树准则2,parent节点一定存在父节点,即祖父节点gparent(根据准则4,其为黑色节点),如果parent位于祖父节点的左侧,进入以下的if处理。

        gparent = rb_red_parent(parent);

        tmp = gparent->rb_right;
        if (parent != tmp) {    /* parent == gparent->rb_left */
2.2.3.1 红色叔节点情况(Case-1)

如果祖父节点的右子节点存在,并且其颜色也是红色,称此情况为Case-1,如代码中的图示。由于插入节点为红色,其父节点也是红色,违反了准则4。如下操作,将其父节点和叔节点都设置为黑色,并且祖父节点设置为红色。

以上改动之后,祖父节点为红色,但是其父节点也有可能为红色,不符合准则4,此时,将祖父节点作为新插入节点(node),其父节点设置为parent,重新执行插入(continue)。另外,新插入节点为其父节点的左子节点还是右子节点情况都是一样的。

以下注释的图示中,小写字母表示红色节点,大写字母表示黑色节点。

            if (tmp && rb_is_red(tmp)) {
                /*
                 * Case 1 - node's uncle is red (color flips).
                 *
                 *       G            g
                 *      / \          / \
                 *     p   u  -->   P   U
                 *    /            /
                 *   n            n
                 *
                 * However, since g's parent might be red, and
                 * 4) does not allow this, we need to recurse at g.
                 */
                rb_set_parent_color(tmp, gparent, RB_BLACK);
                rb_set_parent_color(parent, gparent, RB_BLACK);
                node = gparent;
                parent = rb_parent(node);
                rb_set_parent_color(node, parent, RB_RED);
                continue;
            }
2.2.3.2 黑色叔节点情况-I(Case-2)

第二种情况Case2,要插入节点的叔节点为黑色,并且新节点为其父节点的右子节点,如下图所示。由于父节点p和新插入节点n都为红色,不符合准则4,下面在父节点上进行左旋,之后还是不符合准则4,这时将新插入节点设置为父节点,tmp赋值为父节点的右子节点,由以下的Cas-3进行处理。

            tmp = parent->rb_right;
            if (node == tmp) {
                /*
                 * Case 2 - node's uncle is black and node is
                 * the parent's right child (left rotate at parent).
                 *
                 *      G             G
                 *     / \           / \
                 *    p   U  -->    n   U
                 *     \           /
                 *      n         p
                 *
                 * This still leaves us in violation of 4), the
                 * continuation into Case 3 will fix that.
                 */
                tmp = node->rb_left;
                WRITE_ONCE(parent->rb_right, tmp);
                WRITE_ONCE(node->rb_left, parent);
                if (tmp)
                    rb_set_parent_color(tmp, parent, RB_BLACK);
                rb_set_parent_color(parent, node, RB_RED);
                augment_rotate(parent, node);
                parent = node;
                tmp = node->rb_right;
            }
2.2.3.3 黑色叔节点情况-II(Case-3)

第三种情况Case3,要插入节点的叔节点(Uncle)为黑色,并且新节点为其父节点的左子节点,如下图所示。由于父节点为红色,叔节点为黑色,根据准则5,叔节点必定是叶子节点,否则其所在子树的黑色节点将多于父节点p所在的子树。

在插入红色新节点n之后,p和n都是红色,不符合准则4,在祖父节点G处进行右旋操作,并且进行颜色调整。父节点p和祖父节点g的颜色做翻转,插入操作完成,退出。

            /*
             * Case 3 - node's uncle is black and node is
             * the parent's left child (right rotate at gparent).
             *
             *        G           P
             *       / \         / \
             *      p   U  -->  n   g
             *     /                 \
             *    n                   U
             */
            WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */
            WRITE_ONCE(parent->rb_right, gparent);
            if (tmp)
                rb_set_parent_color(tmp, gparent, RB_BLACK);
            __rb_rotate_set_parents(gparent, parent, root, RB_RED);
            augment_rotate(gparent, parent);
            break;

2.2.4 父节点为红色情况-II

如果父节点为祖父节点gparent的右子节点。

2.2.4.1 红色叔节点(Case-1)

情况一Case-1,如果叔节点存在并且为红色节点,此时进行颜色翻转:祖父节点由黑变红,u和p节点由红变黑,这导致一个问题,就是祖父节点的上一节父节点有可能为红色,违反准则4,所以,需要将祖父节点作为新插入节点,重新执行。

        } else {
            tmp = gparent->rb_left;
            if (tmp && rb_is_red(tmp)) {
                /* Case 1 - color flips */
                 *       G            g
                 *      / \          / \
                 *     u   p  -->   U   P
                 *          \            \
                 *           n            n
                 */
                rb_set_parent_color(tmp, gparent, RB_BLACK);
                rb_set_parent_color(parent, gparent, RB_BLACK);
                node = gparent;
                parent = rb_parent(node);
                rb_set_parent_color(node, parent, RB_RED);
                continue;
            }
2.2.4.2 黑色叔节点-I(Case-2)

情况二Case-2,新插入节点为父节点的左子节点,此时在父节点处执行右旋操作,但是仍然违反准则4。此时,将新插入节点node设置为父节点,而将其左子节点(之前的parent)设置为tmp,由以下的Case-3进行处理。

            tmp = parent->rb_left;
            if (node == tmp) {
                /* Case 2 - right rotate at parent */
                 *        G           G
                 *       / \         / \
                 *      U   p  -->  U   n
                 *         /             \
                 *        n               p
                 */
                tmp = node->rb_right;
                WRITE_ONCE(parent->rb_left, tmp);
                WRITE_ONCE(node->rb_right, parent);
                if (tmp)
                    rb_set_parent_color(tmp, parent, RB_BLACK);
                rb_set_parent_color(parent, node, RB_RED);
                augment_rotate(parent, node);
                parent = node;
                tmp = node->rb_left;
            }
2.2.4.3 黑色叔节点-II(Case-3)

情况三Case-3,新插入节点为父节点的右子节点,此时在祖父节点处执行左旋操作,并且将原来的祖父节点g由黑色设置为红色。插入操作完成,退出。

            /* Case 3 - left rotate at gparent */
             *
             *        G           p
             *       / \         / \
             *      U   p  -->  g   n
             *           \     /      
             *            n   U   
             */
            WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */
            WRITE_ONCE(parent->rb_left, gparent);
            if (tmp)
                rb_set_parent_color(tmp, gparent, RB_BLACK);
            __rb_rotate_set_parents(gparent, parent, root, RB_RED);
            augment_rotate(gparent, parent);
            break;
        }
    }

内核版本 5.0

©️2020 CSDN 皮肤主题: 编程工作室 设计师:CSDN官方博客 返回首页