Redis6.0.6源码阅读——基础数据结构(rax基数树)【上】

1,248 阅读10分钟

前言

Redis5引入了一个Stream,根据开发者的自述Stream借鉴了Kafka的功能,所以相当于在Redis里面内置了一个小型的Kafka。查阅相关文章,发现Stream是基于rax基数树实现的。早前阅读过某Redis书籍,但是作者说rax的源码实现起来很复杂,并没有展开分析,所以看一下rax源码来挑战一下。

正文

rax结构

 *
 *              (f) ""
 *                \
 *                (o) "f"
 *                  \
 *                  (o) "fo"
 *                    \
 *                  [t   b] "foo"
 *                  /     \
 *         "foot" (e)     (a) "foob"
 *                /         \
 *      "foote" (r)         (r) "fooba"
 *              /             \
 *    "footer" []             [] "foobar"
 *
 *
 *                  ["foo"] ""
 *                     |
 *                  [t   b] "foo"
 *                  /     \
 *        "foot" ("er")    ("ar") "foob"
 *                 /          \
 *       "footer" []          [] "foobar"
 
 *
 *
 *                    (f) ""
 *                    /
 *                 (i o) "f"
 *                 /   \
 *    "firs"  ("rst")  (o) "fo"
 *              /        \
 *    "first" []       [t   b] "foo"
 *                     /     \
 *           "foot" ("er")    ("ar") "foob"
 *                    /          \
 *          "footer" []          [] "foobar"
 *
 */

这是代码里面注释出来的三幅图,初步了解到基数树就是一个前缀树,每一个节点上都有一个指针指向下一个节点和数据,但是作者表明对于连续的字符串可以压缩指针,只保留一个指针指向下一个节点(不懂看下图)

typedef struct raxNode {
    uint32_t iskey:1;     //是否包含key 1字节
    uint32_t isnull:1;    //value是否为空 1字节
    uint32_t iscompr:1;   //是否为压缩节点 1字节
    uint32_t size:29;     //如果是压缩节点:保存长度;普通节点保存子节点数量
    unsigned char data[];
} raxNode;

typedef struct rax {
    raxNode *head; //root节点
    uint64_t numele; //元素的数量
    uint64_t numnodes; //节点的数量
} rax;

以上是node和rax的结构体,如下图所示:

node.png

header信息占用32bit4个字节,以此可以计算出内存对齐字节填充量:

#define raxPadding(nodesize) ((sizeof(void*)-((nodesize+4) % sizeof(void*))) & (sizeof(void*)-1))

通过nodeSize计算,sizeof(void*)可以根据当前系统位数不同返回不同,64位=8,32位=4

有了这些信息,也可以计算出节点长度,节点第一个child指针以及最后一个child的指针

#define raxNodeLastChildPtr(n) ((raxNode**) ( \
    ((char*)(n)) + \
    raxNodeCurrentLength(n) - \
    sizeof(raxNode*) - \
    (((n)->iskey && !(n)->isnull) ? sizeof(void*) : 0) \
))

#define raxNodeFirstChildPtr(n) ((raxNode**) ( \
    (n)->data + \
    (n)->size + \
    raxPadding((n)->size)))

#define raxNodeCurrentLength(n) ( \
    sizeof(raxNode)+(n)->size+ \
    raxPadding((n)->size)+ \
    ((n)->iscompr ? sizeof(raxNode*) : sizeof(raxNode*)*(n)->size)+ \
    (((n)->iskey && !(n)->isnull)*sizeof(void*)) \
)

比如想要计算第一个child指针,只需要找到data数据后的第一个内存地址,简单相加,而最后一个child就要通过总长度相减


普通节点.png

压缩节点.png

以上是真实数据普通节点和压缩节点的差别在于压缩节点上只有末尾字符串的指针

++简单模拟一下数据的插入++

一个abc的压缩节点 原abc.png

插入ab.png

image.png

API

rax *raxNew(void) {
    rax *rax = rax_malloc(sizeof(*rax));
    if (rax == NULL) return NULL;
    rax->numele = 0;
    rax->numnodes = 1;
    rax->head = raxNewNode(0,0);
    if (rax->head == NULL) {
        rax_free(rax);
        return NULL;
    } else {
        return rax;
    }
}

创建一个空rax树,head节点是一个空节点,没有数据

raxNode *raxNewNode(size_t children, int datafield) {
    size_t nodesize = sizeof(raxNode)+children+raxPadding(children)+
                      sizeof(raxNode*)*children;
    if (datafield) nodesize += sizeof(void*);
    raxNode *node = rax_malloc(nodesize);
    if (node == NULL) return NULL;
    node->iskey = 0;
    node->isnull = 0;
    node->iscompr = 0;
    node->size = children;
    return node;
}

预先创建一个有n个children的节点,如果是一个数据节点,还要额外分配sizeof(void*)来容差指针

下面是rax树的插入方法,这个方法很长,需要慢慢看

int raxGenericInsert(rax *rax, unsigned char *s, size_t len, void *data, void **old, int overwrite) {
    size_t i;
    //分裂位置
    int j = 0;
    raxNode *h, **parentlink;

    debugf("### Insert %.*s with value %p\n", (int)len, s, data);
    i = raxLowWalk(rax,s,len,&h,&parentlink,&j,NULL);

想要插入一个数据,必须得获取插入的位置,从插入字符串开始匹配,匹配到不相等位置。所以最终需要的数据有:

  • 原字符串匹配的位置:i
  • 匹配到了哪一个节点指针:h
  • 匹配到节点的上数据的位置:j
  • 匹配节点的父节点指针:parentlink

插入获取.png

以上信息如图所示

static inline size_t raxLowWalk(rax *rax, unsigned char *s, size_t len, raxNode **stopnode, raxNode ***plink, int *splitpos, raxStack *ts) {
    raxNode *h = rax->head;
    raxNode **parentlink = &rax->head;

    size_t i = 0; /* 待插入字符串的索引*/
    size_t j = 0; /* 子节点的指针 如果是压缩节点为字节*/
    while(h->size && i < len) {
        debugnode("Lookup current node",h);
        unsigned char *v = h->data;

        if (h->iscompr) {
            //压缩节点则要逐个字节比较
            for (j = 0; j < h->size && i < len; j++, i++) {
                //压缩节点如果遇到不相等表示分裂位置
                if (v[j] != s[i]) break;
            }
            //退出while 数据没有遍历完 不需要跟新其他信息
            if (j != h->size) break;
        } else {
            //相等表示进入下一个子节点
            for (j = 0; j < h->size; j++) {
                if (v[j] == s[i]) break;
            }
            //表示已经完全匹配上了
            if (j == h->size) break;
            i++;
        }

        if (ts) raxStackPush(ts,h); /* Save stack of parent nodes. */
        raxNode **children = raxNodeFirstChildPtr(h);
        if (h->iscompr) j = 0; /* Compressed node only child is at index 0. */
        memcpy(&h,children+j,sizeof(h));
        parentlink = children+j;
        j = 0; /* If the new node is non compressed and we do not
                  iterate again (since i == len) set the split
                  position to 0 to signal this node represents
                  the searched key. */
    }
    debugnode("Lookup stop node is",h);
    if (stopnode) *stopnode = h;
    if (plink) *plink = parentlink;
    if (splitpos && h->iscompr) *splitpos = j;
    return i;
}

如果压缩节点 那么必须得完全匹配才能跳到下一个节点,普通节点则不断匹配和跳转节点,直到不相等位置,然后记录信息。

if (i == len && (!h->iscompr || j == 0 )) {
        debugf("### Insert: node representing key exists\n");
        //如果不是key节点 或者 节点是空的并且支持覆盖
        if (!h->iskey || (h->isnull && overwrite)) {
            h = raxReallocForData(h,data);
            if (h) memcpy(parentlink,&h,sizeof(h));
        }
        //上面过程中如果节点是NULL 说明内存分配不成功 没有足够内存
        if (h == NULL) {
            errno = ENOMEM;
            return 0;
        }

        //如果存在key 支持覆盖则覆盖
        if (h->iskey) {
            if (old) *old = raxGetData(h);
            if (overwrite) raxSetData(h,data);
            errno = 0;
            return 0;
        }

        //插入节点 增加数量
        raxSetData(h,data);
        rax->numele++;
        return 1;
    }

当获取到插入位置后,如果i==len则表示字符串被完全匹配了,字符串被完全匹配且不是压缩节点(是压缩节点也可以,但是j得等于0,说明压缩节点也被完全匹配了),说明不需要进行节点的分裂,可以进行设置data,这一段代码则是进行简单情况的判断,然后完成插入。

其他情况就要开始分裂了,注释上说明了几种情况,能够涵盖插入时的情形。

/*
原rax树
"ANNIBALE" -> "SCO" -> []
*/

     * 1) Inserting "ANNIENTARE"
     *
     *               |B| -> "ALE" -> "SCO" -> []
     *     "ANNI" -> |-|
     *               |E| -> (... continue algo ...) "NTARE" -> []
     *
     * 2) Inserting "ANNIBALI"
     *
     *                  |E| -> "SCO" -> []
     *     "ANNIBAL" -> |-|
     *                  |I| -> (... continue algo ...) []
     *
     * 3) Inserting "AGO" (Like case 1, but set iscompr = 0 into original node)
     *
     *            |N| -> "NIBALE" -> "SCO" -> []
     *     |A| -> |-|
     *            |G| -> (... continue algo ...) |O| -> []
     *
     * 4) Inserting "CIAO"
     *
     *     |A| -> "NNIBALE" -> "SCO" -> []
     *     |-|
     *     |C| -> (... continue algo ...) "IAO" -> []
     *
     * 5) Inserting "ANNI"
     *
     *     "ANNI" -> "BALE" -> "SCO" -> []
     *

它将这5中情况分为两类算法:

  • 1-4情况算法:插入的数据只被匹配了一点
    • 保存NEXT指针:对应上图SCO数据
    • 如果 j为 0,表示匹配到压缩节点的开始,可以直接拷贝数据到分裂的节点上。
    • 如果 j 不为 0,裁剪分裂节点,只需要j个字符。
    • 如果分裂后压缩节点剩余的字符长度不为 0,将剩余的节点截断创建新节点,然后指向SCO**
    • 如果 分裂后剩余 节点长度为 0,则直接使用 SCO
  • 5情况算法:插入的数据被完全匹配了,截断。

ANNIBALE" -> "SCO" -> [],为原rax树,来插入情况3AGO数据来试试

首先匹配到了“A”字符串,j=1,i=1。这属于算法1,只匹配了原字符串部分的情况

if (h->iscompr && i != len) {
        debugf("ALGO 1: Stopped at compressed node %.*s (%p)\n",
            h->size, h->data, (void*)h);
        debugf("Still to insert: %.*s\n", (int)(len-i), s+i);
        debugf("Splitting at %d: '%c'\n", j, ((char*)h->data)[j]);
        debugf("Other (key) letter is '%c'\n", s[i]);

        //获取 孩子指针 并且将这个指针上内存拷贝
        raxNode **childfield = raxNodeLastChildPtr(h);
        raxNode *next;
        memcpy(&next,childfield,sizeof(next));
        debugf("Next is %p\n", (void*)next);
        debugf("iskey %d\n", h->iskey);
        if (h->iskey) {
            debugf("key value is %p\n", raxGetData(h));
        }

首先获取next指针,指针上数据为SCO,是一个单独的节点,这个节点只用于连接

 //保存节点当前匹配到的长度 以及还剩多长
        size_t trimmedlen = j;
        size_t postfixlen = h->size - j - 1;
        int split_node_is_key = !trimmedlen && h->iskey && !h->isnull;
        size_t nodesize;

        //创建一个拆分节点
        raxNode *splitnode = raxNewNode(1, split_node_is_key);
        raxNode *trimmed = NULL;
        raxNode *postfix = NULL;

        //分配拆分前后的内存
        if (trimmedlen) {
            nodesize = sizeof(raxNode)+trimmedlen+raxPadding(trimmedlen)+
                       sizeof(raxNode*);
            if (h->iskey && !h->isnull) nodesize += sizeof(void*);
            trimmed = rax_malloc(nodesize);
        }

        if (postfixlen) {
            nodesize = sizeof(raxNode)+postfixlen+raxPadding(postfixlen)+
                       sizeof(raxNode*);
            postfix = rax_malloc(nodesize);
        }
	/* 内存不够中断 并且释放 */
        if (splitnode == NULL ||
            (trimmedlen && trimmed == NULL) ||
            (postfixlen && postfix == NULL))
        {
            rax_free(splitnode);
            rax_free(trimmed);
            rax_free(postfix);
            errno = ENOMEM;
            return 0;
        }
	splitnode->data[0] = h->data[j];

然后要把“ANNIBALE”给拆开,创建一个长度为1的splitnode,splitnode用于存放“A”,分配拆分前后的内存。

//j等于0可以直接copy整个节点数据
        if (j == 0) {
            /* 3a: Replace the old node with the split node. */
            if (h->iskey) {
                void *ndata = raxGetData(h);
                raxSetData(splitnode,ndata);
            }
            memcpy(parentlink,&splitnode,sizeof(splitnode));

当j==0的时候可以直接设置data,但是现在不属于该情况

//裁剪压缩节点
            trimmed->size = j;
            memcpy(trimmed->data,h->data,j);
            trimmed->iscompr = j > 1 ? 1 : 0;
            trimmed->iskey = h->iskey;
            trimmed->isnull = h->isnull;
            if (h->iskey && !h->isnull) {
                void *ndata = raxGetData(h);
                raxSetData(trimmed,ndata);
            }
            raxNode **cp = raxNodeLastChildPtr(trimmed);
            memcpy(cp,&splitnode,sizeof(splitnode));
            memcpy(parentlink,&trimmed,sizeof(trimmed));
            parentlink = cp;
            rax->numnodes++;
        }

然后开始裁剪“N”,从N开始不匹配了,将A和N相连,并且跟新parentlink,这里开始已经分裂出一个节点了,所以numnodes增加

//提取拆分后剩下节点
        if (postfixlen) {
            postfix->iskey = 0;
            postfix->isnull = 0;
            postfix->size = postfixlen;
            postfix->iscompr = postfixlen > 1;
            memcpy(postfix->data,h->data+j+1,postfixlen);
            raxNode **cp = raxNodeLastChildPtr(postfix);
            memcpy(cp,&next,sizeof(next));
            rax->numnodes++;
        } else {
            postfix = next;
        }
        
        raxNode **splitchild = raxNodeLastChildPtr(splitnode);
        memcpy(splitchild,&postfix,sizeof(postfix));
        
        rax_free(h);
        h = splitnode;

这个时候还剩下"NIBALE",将这段和“SCO”相连,最后再将”N“和”NIBALE“相连,原来的h节点可以释放,并且将节点指向splitnode(N)

到现在压缩节点分裂完成,但是还没有插入新节点,别着急,看一下算法2代码

else if (h->iscompr && i == len) {
    /* ------------------------- ALGORITHM 2 --------------------------- */
        debugf("ALGO 2: Stopped at compressed node %.*s (%p) j = %d\n",
            h->size, h->data, (void*)h, j);

        //将原压缩节点分为两个 :trimmed postfix
        size_t postfixlen = h->size - j;
        size_t nodesize = sizeof(raxNode)+postfixlen+raxPadding(postfixlen)+
                          sizeof(raxNode*);
        if (data != NULL) nodesize += sizeof(void*);
        raxNode *postfix = rax_malloc(nodesize);

        nodesize = sizeof(raxNode)+j+raxPadding(j)+sizeof(raxNode*);
        if (h->iskey && !h->isnull) nodesize += sizeof(void*);
        raxNode *trimmed = rax_malloc(nodesize);

        if (postfix == NULL || trimmed == NULL) {
            rax_free(postfix);
            rax_free(trimmed);
            errno = ENOMEM;
            return 0;
        }

        //保存压缩节点的子节点 最终连接情况 trimmed —— postfix —— next
        raxNode **childfield = raxNodeLastChildPtr(h);
        raxNode *next;
        memcpy(&next,childfield,sizeof(next));


        postfix->size = postfixlen;
        postfix->iscompr = postfixlen > 1;
        postfix->iskey = 1;
        postfix->isnull = 0;
        memcpy(postfix->data,h->data+j,postfixlen);
        raxSetData(postfix,data);
        raxNode **cp = raxNodeLastChildPtr(postfix);
        //连接 postfix —— next
        memcpy(cp,&next,sizeof(next));
        rax->numnodes++;

        trimmed->size = j;
        trimmed->iscompr = j > 1;
        trimmed->iskey = 0;
        trimmed->isnull = 0;
        memcpy(trimmed->data,h->data,j);
        //连接parentlink —— trimmed
        memcpy(parentlink,&trimmed,sizeof(trimmed));
        if (h->iskey) {
            void *aux = raxGetData(h);
            raxSetData(trimmed,aux);
        }

        cp = raxNodeLastChildPtr(trimmed);
        //连接 trimmed - postfix
        memcpy(cp,&postfix,sizeof(postfix));

        rax->numele++;
        rax_free(h);
        return 1;
    }

这个算法很简单,看注释

   //到了这一步已经将原来的rax树分裂了 但是还没有插入新节点
    while(i < len) {
        raxNode *child;

        //首先会进入else判断 因为h上面size是大于0的 h为spiltnode 
        if (h->size == 0 && len-i > 1) {
            debugf("Inserting compressed node\n");
            size_t comprsize = len-i;
            if (comprsize > RAX_NODE_MAX_SIZE)
                comprsize = RAX_NODE_MAX_SIZE;
            raxNode *newh = raxCompressNode(h,s+i,comprsize,&child);
            if (newh == NULL) goto oom;
            h = newh;
            memcpy(parentlink,&h,sizeof(h));
            parentlink = raxNodeLastChildPtr(h);
            i += comprsize;
        } else {
            //会在h插入 s的一个字符串 并且修改h 这个时候就会进入上面if判断
            debugf("Inserting normal node\n");
            raxNode **new_parentlink;
            raxNode *newh = raxAddChild(h,s[i],&child,&new_parentlink);
            if (newh == NULL) goto oom;
            h = newh;
            memcpy(parentlink,&h,sizeof(h));
            parentlink = new_parentlink;
            i++;
        }
        rax->numnodes++;
        h = child;
    }
    raxNode *newh = raxReallocForData(h,data);
    if (newh == NULL) goto oom;
    h = newh;
    if (!h->iskey) rax->numele++;
    raxSetData(h,data);
    memcpy(parentlink,&h,sizeof(h));
    return 1;

在这个阶段 rax树为|A|-|N| -> "NIBALE" -> "SCO" -> []

h是 “N”

首先进入else判断,在“N”后面插入”G“,然后修改h为”G“的子节点,这个时候“G”是没有字节点的,进入if判断吧将后面的字符串包装成压缩节点插入。


总结:插入算法是真的复杂。。。。。

我们可以把一个压缩节点上的字符串和插入的字符串分为三个部分:公共前缀、前缀后第一个字符、剩余字符

公共前缀可以不用管它,前缀后的第一个字符串拼凑成一起,剩余后缀字符形成两条分支。

算法1

插入总结.png

算法2

插入总结2.png

结尾

本文讲述了rax基数树的数据结构以及指针计算方式,分析了插入算法的源码,通过两种算法配合能够适应不同情况的插入,下一篇文章将剩下的方法源码分析一下。