定时器实现之红黑树（二）

1.概述

        书接上回定时器实现之最小堆（一），今天采用红黑树来作为定时器的容器，组织定时器的定时节点。

2.为什么红黑树能用来实现定时器

         前面一章提到过，能用来实现定时器的容器的基本要求就是有序，而红黑树的中序遍历就是有序的，如下图：

        并且，对于红黑树中最小和最大的节点分别位于最左和最右子树，基于这一现象，我们在定时器的实现中，对于找最近要过期的任务，只需要最左子树的叶子节点就行，这是我们高效判定定时任务是否触发的核心。

3.代码实现 

        对于红黑树用于定时器的实现参考了nginx 并增加了一些自己的东西具体如下：

//rbtree.h
#pragma once

struct rbtree_node {
    unsigned long long key;
    rbtree_node*       left;
    rbtree_node*       right;
    rbtree_node*       parent;
    unsigned char      color; // 0: red, 1: black
};

struct rbtree {
    rbtree_node* root;
    rbtree_node* sentinel; //哨兵节点
};

void rbtree_init(rbtree* r, rbtree_node* s);

void rbtree_insert(rbtree *tree, rbtree_node *node);

void rbtree_delete(rbtree *tree, rbtree_node *node);

rbtree_node* rbtree_min(rbtree *tree);

rbtree_node* rbtree_max(rbtree *tree);

//rbtree.cpp

#include "rbtree.h"


static void insert(rbtree_node* r, rbtree_node* n, rbtree_node* s) {
    rbtree_node  **p = nullptr;

    for ( ;; ) {
        p = (n->key < r->key) ? &r->left : &r->right;

        if (*p == s) {
            break;
        }

        r = *p;
    }

    *p = n;
    n->parent = r;
    n->left = s;
    n->right = s;
    n->color = 1;
}

static void rbtree_left_rotate(rbtree_node**r, rbtree_node* s, rbtree_node* n) {
    rbtree_node* temp = n->right;
    n->right = n->right->left;

    if(temp->left != s) {
        temp->left->parent = n;
    }

    temp->parent = n->parent;

    if(n == *r) {
        *r = temp;
    } else if(n == n->parent->left) {
        n->parent->left = temp;
    } else {
        n->parent->right = temp;
    }

    temp->left = n;
    n->parent = temp;
}

static void rbtree_right_rotate(rbtree_node**r, rbtree_node* s, rbtree_node* n) {
    rbtree_node  *temp;

    temp = n->left;
    n->left = temp->right;

    if (temp->right != s) {
        temp->right->parent = n;
    }

    temp->parent = n->parent;

    if (n == *r) {
        *r = temp;

    } else if (n == n->parent->right) {
        n->parent->right = temp;

    } else {
        n->parent->left = temp;
    }

    temp->right = n;
    n->parent = temp;
}

static rbtree_node* rbtree_min(rbtree_node* n, rbtree_node* s) {
    if(!n) {
        return nullptr;
    }

    while(n->left != s) {
        n = n->left;
    }

    return n;
}

static rbtree_node* rbtree_max(rbtree_node* n, rbtree_node* s) {
    if(!n) {
        return nullptr;
    }

    while(n->right != s) {
        n = n->right;
    }

    return n;
}

void rbtree_init(rbtree* r, rbtree_node* s) {
    s->color = 0;
    r->root = s;
    r->sentinel = s;
}

void rbtree_insert(rbtree *tree, rbtree_node *node) {
    rbtree_node** root = &tree->root;
    rbtree_node* sentinel = tree->sentinel;

    if(*root == sentinel) {
        node->parent = nullptr;
        node->left = sentinel;
        node->right = sentinel;
        node->color = 0;
        *root = node;
        return;
    }

    // 插入节点
    insert(*root, node, sentinel);

    rbtree_node* temp = nullptr;
    //判断是否需要调整使其平衡
    while(node != *root && node->parent->color) {
        if(node->parent == node->parent->parent->left) {
            temp = node->parent->parent->right;
            //如果叔父节点是红色则直接变色并进入下一轮调整
            if(temp->color) {
                node->parent->color = 0;
                temp->color = 0;
                node->parent->parent->color = 1;
                node = node->parent->parent;
            } else {
                //如果当前插入节点是父节点的左节点先对父节点进行左旋    
                if(node == node->parent->right) {
                    node = node->parent;
                    rbtree_left_rotate(root, sentinel, node);
                }

                node->parent->color = 0;
                node->parent->parent->color = 1;
                rbtree_right_rotate(root, sentinel, node->parent->parent);
            }
        } else {
            temp = node->parent->parent->left;
            //如果叔父节点是红色则直接变色并进入下一轮调整
            if(temp->color) {
                node->parent->color = 0;
                temp->color = 0;
                node->parent->parent->color = 1;
                node = node->parent->parent;
            } else {
                //如果当前插入节点是父节点的左节点先对父节点进行右旋
                if(node == node->parent->left) {
                    node = node->parent;
                    rbtree_right_rotate(root, sentinel, node);
                }

                node->parent->color = 0;
                node->parent->parent->color = 1;
                rbtree_left_rotate(root, sentinel, node->parent->parent);
            }
        }
    }

    (*root)->color = 0;
}

void rbtree_delete(rbtree *tree, rbtree_node *node) {
    rbtree_node** root = &tree->root;
    rbtree_node* sentinel = tree->sentinel;
    rbtree_node* subst = nullptr;
    rbtree_node* temp = nullptr;

    if (node->left == sentinel) {
        temp = node->right;
        subst = node;

    } else if (node->right == sentinel) {
        temp = node->left;
        subst = node;
    } else {
        subst = rbtree_min(node->right, sentinel);
        if (subst->left != sentinel) {
            temp = subst->left;
        } else {
            temp = subst->right;
        }
    }

    if (subst == *root) {
        *root = temp;
        temp->color = 0;
        return;
    }

    if (subst == subst->parent->left) {
        subst->parent->left = temp;
    } else {
        subst->parent->right = temp;
    }

    if (subst == node) {
        temp->parent = subst->parent;
    } else {
        if (subst->parent == node) {
            temp->parent = subst;

        } else {
            temp->parent = subst->parent;
        }

        subst->left = node->left;
        subst->right = node->right;
        subst->parent = node->parent;
        subst->color = node->color;

        if (node == *root) {
            *root = subst;
        } else {
            if (node == node->parent->left) {
                node->parent->left = subst;
            } else {
                node->parent->right = subst;
            }
        }

        if (subst->left != sentinel) {
            subst->left->parent = subst;
        }

        if (subst->right != sentinel) {
            subst->right->parent = subst;
        }
    }

    if (subst->color) {
        return;
    }

    rbtree_node* w = nullptr;
    while (temp != *root && !temp->color) {

        if (temp == temp->parent->left) {
            w = temp->parent->right;

            if (w->color) {
                w->color = 1;
                temp->parent->color = 0;
                rbtree_left_rotate(root, sentinel, temp->parent);
                w = temp->parent->right;
            }

            if (!w->left->color && !w->right->color) {
                w->color = 0;
                temp = temp->parent;

            } else {
                if (!w->right->color) {
                    w->left->color = 0;
                    w->color = 1;
                    rbtree_right_rotate(root, sentinel, w);
                    w = temp->parent->right;
                }

                w->color = temp->color;
                temp->parent->color = 0;
                w->right->color = 0;
                rbtree_left_rotate(root, sentinel, temp->parent);
                temp = *root;
            }

        } else {
            w = temp->parent->left;

            if (w->color) {
                w->color = 0;
                temp->parent->color = 1;
                rbtree_right_rotate(root, sentinel, temp->parent);
                w = temp->parent->left;
            }

            if (!w->left->color && !w->right->color) {
                w->color = 1;
                temp = temp->parent;

            } else {
                if (!w->left->color) {
                    w->right->color = 0;
                    w->color = 1;
                    rbtree_left_rotate(root, sentinel, w);
                    w = temp->parent->left;
                }

                w->color = temp->parent->color;
                temp->parent = 0;
                w->left->color = 0;
                rbtree_right_rotate(root, sentinel, temp->parent);
                temp = *root;
            }
        }
    }

    temp->color = 0;
}

rbtree_node* rbtree_min(rbtree *tree) {
    if(!tree->root || tree->root == tree->sentinel) {
        return nullptr;
    }

    return rbtree_min(tree->root, tree->sentinel);
}

rbtree_node* rbtree_max(rbtree *tree) {
    if(!tree->root || tree->root == tree->sentinel) {
        return nullptr;
    }

    return rbtree_max(tree->root, tree->sentinel);
}

rbtree_node* ngx_rbtree_next(rbtree *tree, rbtree_node *node)
{
    rbtree_node* root = nullptr;
    rbtree_node* sentinel = nullptr;
    rbtree_node* parent = nullptr;

    sentinel = tree->sentinel;

    if (node->right != sentinel) {
        return rbtree_min(node->right, sentinel);
    }

    root = tree->root;

    for ( ;; ) {
        parent = node->parent;

        if (node == root) {
            return nullptr;
        }

        if (node == parent->left) {
            return parent;
        }

        node = parent;
    }
}

        再对之前的main.cpp进行一些更改：

//main.cpp

#include "minheap.h"
#include "rbtree.h"

#include <cstdint>
#include <iostream>
#include <functional>
#include <chrono>
#include <ctime>
#include <thread>

enum TimeType {
    MIN_HEAP,
    RBTREE,
};

extern "C"{
#define offsetofs(s,m) (size_t)(&reinterpret_cast<const volatile char&>((((s*)0)->m)))
#define container_of(ptr, type, member) ({                  \
    const typeof( ((type *)0)->member ) *__mptr = (ptr);    \
    (type *)( (char *)__mptr - offsetofs(type, member) );})
}

using CallBack = std::function<void(void)>;

struct TimeNode {
    timer_entry env;
    CallBack cb;
};

struct TimeNodeRb {
    rbtree_node env;
    CallBack cb;
};

class Timer {
public:
    Timer(uint32_t size);
    Timer(uint32_t size, TimeType Type);
    Timer() = delete;
    ~Timer();

    int addTimer(uint64_t time, CallBack cb);

    void run();

    void stop();

private:
    uint64_t GetNowTime();

    void AddMinHeapTimer(uint64_t time, CallBack cb);

    void AddRbtreeTimer(uint64_t time, CallBack cb);
private:
    min_heap* _heap;
    rbtree*   _rbtree;
    bool _close;
    TimeType _Type;
};

Timer::Timer(uint32_t size) {
    min_heap_init(&_heap, size);
    _close = false;
}

Timer::Timer(uint32_t size, TimeType Type) {
    switch (Type) {
    case  MIN_HEAP:
        min_heap_init(&_heap, size);
        break;
    case RBTREE:
        _rbtree = new rbtree();
        if(_rbtree) {
            _rbtree->sentinel = new rbtree_node();
            rbtree_init(_rbtree, _rbtree->sentinel);
        }
        break;
    }
    _close = false;
    _Type = Type;
}

Timer::~Timer() {
    switch (_Type) {
    case MIN_HEAP:
        min_heap_free(_heap);
        break;
    case RBTREE:
        delete _rbtree->sentinel;
        delete _rbtree;      
        break;
    }
}

int Timer::addTimer(uint64_t time, CallBack cb) {
    switch (_Type) {
    case MIN_HEAP:
        AddMinHeapTimer(time, cb);
        break;
    case RBTREE:
        AddRbtreeTimer(time, cb);
        break;
    }
    return 0;
}

uint64_t Timer::GetNowTime() {
    auto now = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch());
    uint64_t timestamp_ms = duration.count();
    return timestamp_ms;
}

void Timer::run() {
    while(!_close) {
        uint64_t sleep = 50;
        uint64_t now = GetNowTime();
        switch (_Type) {
            case MIN_HEAP: {
                timer_entry *entry = nullptr;
                if (!min_heap_top(_heap, &entry)) {
                    if (entry->time <= now) {
                        if (!min_heap_pop(_heap, &entry)) {
                            TimeNode *node = container_of(entry, TimeNode, env);
                            if (node) {
                                node->cb();
                            }

                            delete node;
                        }
                    } else {
                        sleep = entry->time - now;
                    }
                }
            }
            case RBTREE: {
                rbtree_node* node = rbtree_min(_rbtree);
                if(!node) {
                    break;
                }

                if(node->key <= now) {
                    rbtree_delete(_rbtree, node);
                    TimeNodeRb* n = container_of(node, TimeNodeRb, env);
                    n->cb();

                    delete n;
                } else {
                    sleep = node->key - now; 
                }
            }
        }

        std::this_thread::sleep_for(std::chrono::milliseconds(sleep));
    }
}


void Timer::stop() {
    _close = true;
}

void Timer::AddMinHeapTimer(uint64_t time, CallBack cb) {
    TimeNode* node = new TimeNode();
    node->env.time = GetNowTime() + time;
    node->cb = cb;
    min_heap_push(_heap, &node->env);
}

void Timer::AddRbtreeTimer(uint64_t time, CallBack cb) {
    TimeNodeRb* node = new TimeNodeRb();
    node->env.key = GetNowTime() + time;
    node->cb = cb;
    rbtree_insert(_rbtree, &node->env);
}

int main(int, char**){

    Timer timer(0, RBTREE);

    auto now = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch());
    uint64_t timestamp_ms = duration.count();
    std::cout << "Start, now time :" << timestamp_ms << std::endl;
    timer.addTimer(1000, [&](void){
        auto now = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch());
        uint64_t timestamp_ms = duration.count();
        std::cout << "timer 1, now time :" << timestamp_ms << std::endl;
    });

    timer.addTimer(2000, [&](void){
        auto now = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch());
        uint64_t timestamp_ms = duration.count();
        std::cout << "timer 2, now time :" << timestamp_ms << std::endl;
    });

    std::thread t([&](){
        timer.run();
    });
    std::this_thread::sleep_for(std::chrono::milliseconds(3000));
    timer.stop();
    t.join();

    return 0;
}

 4.运行效果

         

        整个工程放在Github 上，感兴趣可以直接clone下来用构建运行。

学习参考：0voice · GitHub