高效实时数据排行榜实现

package rank

import (
    "time"

    "common/zebra"
    "logic/db"
    "logic/service/game/rank/zset"
    "logic/service/global"
    "protos/in/db_data"
    "protos/in/r2l"

    "github.com/golang/protobuf/proto"
    l4g "github.com/ivanabc/log4go"
)

// Rank r
type Rank struct {
    redisKey      string
    set           *zset.ZSet
    maxCount      uint32
    changedDB     map[uint64]*zset.ZSkipListNode
    rangeIDRet    []uint64
    rangeScoreRet []uint32
}

// NewRank n
func NewRank(redisKey string, maxCount uint32) *Rank {
    return &Rank{
        redisKey:      redisKey,
        set:           zset.NewZSet(),
        maxCount:      maxCount,
        changedDB:     make(map[uint64]*zset.ZSkipListNode),
        rangeIDRet:    make([]uint64, 0, 5000),
        rangeScoreRet: make([]uint32, 0, 5000),
    }
}

// GetPlayerRank 根据玩家id获取玩家数据
// return - (排名，积分)
func (r *Rank) GetPlayerRank(id uint64) (uint32, uint32) {
    return r.set.Rank(id, true)
}

// LoadRankSync 同步加载排行榜数据库
func (r *Rank) LoadRankSync() {
    if global.StorageMgr == nil {
        return
    }
    items := global.StorageMgr.LoadHashDataSync(r.redisKey)
    if items == nil {
        return
    }
    data := &db_data.RankItemData{}
    for _, v := range items {
        if err := proto.Unmarshal(v, data); err != nil {
            l4g.Error("db loadRank unmarshal error: %s", err.Error())
            continue
        }
        r.InitAdd(data.GetID(), data.GetValue(), data.GetTimeStamp())
    }
}

// InitAdd i
func (r *Rank) InitAdd(id uint64, score uint32, t int64) {
    if r.set.Length() >= r.maxCount && score < r.set.MinScore() {
        r.changedDB[id] = nil
        return
    }
    r.set.Add(score, id, t)
    if r.set.Length() > r.maxCount {
        if ele := r.set.DeleteFirst(); ele != nil {
            r.changedDB[ele.Key()] = nil
        }
    }
}

// ChangeScore c
func (r *Rank) ChangeScore(id uint64, score uint32) bool {
    if r.set.Length() >= r.maxCount && score <= r.set.MinScore() {
        return false
    }

    ele := r.set.Add(score, id, time.Now().UnixNano())
    if ele == nil {
        return false
    }
    r.changedDB[id] = ele
    if r.set.Length() > r.maxCount {
        if ele := r.set.DeleteFirst(); ele != nil {
            r.changedDB[ele.Key()] = nil
        }
    }
    return true
}

// GetRange 取排名[rankBegin, rankEnd]间所有
func (r *Rank) GetRange(rankBegin uint32, rankEnd uint32) ([]uint64, []uint32) {
    r.rangeScoreRet = r.rangeScoreRet[:0]
    r.rangeIDRet = r.rangeIDRet[:0]
    r.set.Range(rankBegin, rankEnd, true, &r.rangeIDRet, &r.rangeScoreRet)
    return r.rangeIDRet, r.rangeScoreRet
}

// GetFirst 获得排行第一
func (r *Rank) GetFirst() *zset.Element {
    return r.set.Tail()
}

// Save 保存数据库
func (r *Rank) Save() {
    if len(r.changedDB) == 0 {
        return
    }
    msg := &db.RedisRequest{
        PH: &zebra.PackHead{
            Cmd: uint32(r2l.ID_MSG_L2R_SaveRank),
        },
    }
    data := &r2l.L2R_SaveRankData{
        Name: r.redisKey,
    }
    for k, v := range r.changedDB {
        if v == nil {
            data.DeleteItems = append(data.DeleteItems, k)
        } else {
            data.Items = append(data.Items, &db_data.RankItemData{
                ID:        v.Element().Key(),
                Value:     v.Score(),
                TimeStamp: v.Element().Time(),
            })
        }
    }
    msg.Data = data
    if global.StorageMgr != nil {
        global.StorageMgr.SendDataToDB(msg)
    }

    r.changedDB = make(map[uint64]*zset.ZSkipListNode)
}

package zset

import (
    "math/rand"
)

const (
    skipListMaxLevel = 8    // (1/p)^maxLevel >= maxNode
    skipListP        = 0.25 // SkipList P = 1/4
)

// Element e
type Element struct {
    time int64
    key  uint64
}

// Key return key
func (e *Element) Key() uint64 {
    return e.key
}

// Time 时间
func (e *Element) Time() int64 {
    return e.time
}

type zSkipListLevel struct {
    forward *ZSkipListNode
    span    uint32
}

// ZSkipListNode is an element of a skip list
type ZSkipListNode struct {
    ele      *Element
    score    uint32
    backward *ZSkipListNode
    level    []zSkipListLevel
    order    int
}

func zslCreateNode(level int, score uint32, ele *Element) *ZSkipListNode {
    zn := &ZSkipListNode{
        ele:   ele,
        score: score,
        level: make([]zSkipListLevel, level),
    }
    return zn
}

// Score return score
func (node *ZSkipListNode) Score() uint32 {
    return node.score
}

// Element return Element
func (node *ZSkipListNode) Element() *Element {
    return node.ele
}

// zSkipList represents a skip list
type zSkipList struct {
    header, tail *ZSkipListNode
    length       uint32
    level        int // current level count
}

// zslCreate creates a skip list
func zslCreate() *zSkipList {
    zsl := &zSkipList{
        level: 1,
    }
    zsl.header = zslCreateNode(skipListMaxLevel, 0, nil)
    return zsl
}

// insert element
func (list *zSkipList) insert(node *ZSkipListNode) *ZSkipListNode {
    var update [skipListMaxLevel]*ZSkipListNode
    var rank [skipListMaxLevel]uint32

    x := list.header
    for i := list.level - 1; i >= 0; i-- {
        if i == list.level-1 {
            rank[i] = 0
        } else {
            rank[i] = rank[i+1]
        }

        for x.level[i].forward != nil &&
            (x.level[i].forward.score < node.score ||
                x.level[i].forward.score == node.score &&
                    node.ele.Time() < x.level[i].forward.ele.Time()) {
            rank[i] += x.level[i].span
            x = x.level[i].forward
        }
        update[i] = x
    }

    level := len(node.level)
    if level > list.level {
        for i := list.level; i < level; i++ {
            rank[i] = 0
            update[i] = list.header
            update[i].level[i].span = list.length
        }
        list.level = level
    }

    x = node
    for i := 0; i < level; i++ {
        x.level[i].forward = update[i].level[i].forward
        update[i].level[i].forward = x
        x.level[i].span = update[i].level[i].span - (rank[0] - rank[i])
        update[i].level[i].span = (rank[0] - rank[i]) + 1
    }
    for i := level; i < list.level; i++ {
        update[i].level[i].span++
    }
    next := x.level[0].forward
    if next != nil && x.score == next.score && x.ele.Time() == next.ele.Time() {
        x.order = next.order + 1
    }

    if update[0] == list.header {
        x.backward = nil
    } else {
        x.backward = update[0]
    }
    if x.level[0].forward == nil {
        list.tail = x
    } else {
        x.level[0].forward.backward = x
    }
    list.length++
    return x
}

// delete element
func (list *zSkipList) delete(node *ZSkipListNode) *ZSkipListNode {
    var update [skipListMaxLevel]*ZSkipListNode
    x := list.header
    for i := list.level - 1; i >= 0; i-- {
        for next := x.level[i].forward; next != nil &&
            (next.score < node.score ||
                next.score == node.score &&
                    (node.ele.Time() < next.ele.Time() ||
                        node.ele.Time() == next.ele.Time() && node.order < next.order)); next = x.level[i].forward {
            x = next
        }
        update[i] = x
    }
    x = x.level[0].forward
    if x != nil && x.score == node.score && x.ele.key == node.ele.key {
        for i := 0; i < list.level; i++ {
            if update[i].level[i].forward == x {
                update[i].level[i].span += x.level[i].span - 1
                update[i].level[i].forward = x.level[i].forward
            } else {
                update[i].level[i].span--
            }
        }
        if x.level[0].forward == nil {
            list.tail = x.backward
        } else {
            x.level[0].forward.backward = x.backward
        }
        for list.level > 1 && list.header.level[list.level-1].forward == nil {
            list.level--
        }

        list.length--
        return x
    }
    return nil
}

// Find the rank for an element.
// Returns 0 when the element cannot be found, rank otherwise.
// Note that the rank is 1-based
func (list *zSkipList) zslGetRank(node *ZSkipListNode) uint32 {
    var rank uint32
    x := list.header
    for i := list.level - 1; i >= 0; i-- {
        for next := x.level[i].forward; next != nil &&
            (next.score < node.score ||
                next.score == node.score &&
                    (node.ele.time < next.ele.time ||
                        node.ele.time == next.ele.time && node.order <= next.order)); next = x.level[i].forward {
            rank += x.level[i].span
            x = next
        }
        if x.ele != nil && x.ele.key == node.ele.key {
            return rank
        }
    }
    return 0
}

func (list *zSkipList) randomLevel() int {
    lvl := 1
    for lvl < skipListMaxLevel && rand.Float64() < skipListP {
        lvl++
    }
    return lvl
}

// Finds an element by its rank. The rank argument needs to be 1-based.
func (list *zSkipList) getElementByRank(rank uint32) *ZSkipListNode {
    if rank == list.length {
        return list.tail
    }

    if rank == 1 {
        return list.header.level[0].forward
    }

    var traversed uint32
    x := list.header
    for i := list.level - 1; i >= 0; i-- {
        for x.level[i].forward != nil && traversed+x.level[i].span <= rank {
            traversed += x.level[i].span
            x = x.level[i].forward
        }
        if traversed == rank {
            return x
        }
    }
    return nil
}

// ZSet set
type ZSet struct {
    dict map[uint64]*ZSkipListNode
    zsl  *zSkipList
}

// NewZSet create ZSet
func NewZSet() *ZSet {
    zs := &ZSet{
        dict: make(map[uint64]*ZSkipListNode),
        zsl:  zslCreate(),
    }
    return zs
}

// Add a new element or update the score of an existing element
func (zs *ZSet) Add(score uint32, key uint64, t int64) *ZSkipListNode {
    if node := zs.dict[key]; node != nil {
        oldScore := node.score
        if score == oldScore {
            return nil
        }
        if next := node.level[0].forward; score > oldScore && (next == nil || score < next.score) {
            node.score = score
            node.ele.time = t
        }  else if score < oldScore && (node.backward == nil || score > node.backward.score) {
            node.score = score
            node.ele.time = t
        } else {
            zs.zsl.delete(node)
            node.score = score
            node.ele.time = t
            zs.zsl.insert(node)
        }
        return node
    } else {
        ele := &Element{
            key:  key,
            time: t,
        }
        lvl := zs.zsl.randomLevel()
        node := zslCreateNode(lvl, score, ele)
        zs.zsl.insert(node)
        zs.dict[key] = node
        return node
    }
}

// Delete the element ‘ele‘ from the sorted set,
// return 1 if the element existed and was deleted, 0 otherwise
func (zs *ZSet) Delete(id uint64) int {
    node := zs.dict[id]
    if node == nil {
        return 0
    }
    zs.zsl.delete(node)
    delete(zs.dict, id)
    return 1
}

// Rank return 1-based rank or 0 if not exist
func (zs *ZSet) Rank(id uint64, reverse bool) (uint32, uint32) {
    node := zs.dict[id]
    if node != nil {
        rank := zs.zsl.zslGetRank(node)
        if rank > 0 {
            if reverse {
                return zs.zsl.length - rank + 1, node.score
            }
            return rank, node.score
        }
    }
    return 0, 0
}

// Score return score
func (zs *ZSet) Score(id uint64) uint32 {
    node := zs.dict[id]
    if node != nil {
        return node.score
    }
    return 0
}

// Range return 1-based elements in [start, end]
func (zs *ZSet) Range(start uint32, end uint32, reverse bool, retKey *[]uint64, retScore *[]uint32) {
    if start == 0 {
        start = 1
    }
    if end == 0 {
        end = zs.zsl.length
    }
    if start > end || start > zs.zsl.length {
        return
    }
    if end > zs.zsl.length {
        end = zs.zsl.length
    }
    rangeLen := end - start + 1
    if reverse {
        node := zs.zsl.getElementByRank(zs.zsl.length - start + 1)
        for i := uint32(0); i < rangeLen; i++ {
            *retKey = append(*retKey, node.ele.key)
            *retScore = append(*retScore, node.score)
            node = node.backward
        }
    } else {
        node := zs.zsl.getElementByRank(start)
        for i := uint32(0); i < rangeLen; i++ {
            *retKey = append(*retKey, node.ele.key)
            *retScore = append(*retScore, node.score)
            node = node.level[0].forward
        }
    }
}

// Length return the element count
func (zs *ZSet) Length() uint32 {
    return zs.zsl.length
}

// MinScore return min score
func (zs *ZSet) MinScore() uint32 {
    first := zs.zsl.header.level[0].forward
    if first != nil {
        return first.score
    }
    return 0
}

// Tail return the last element
func (zs *ZSet) Tail() *Element {
    if zs.zsl.tail != nil {
        return zs.zsl.tail.ele
    }
    return nil
}

// DeleteFirst the first element
func (zs *ZSet) DeleteFirst() *Element {
    node := zs.zsl.header.level[0].forward
    zs.zsl.delete(node)
    delete(zs.dict, node.ele.key)
    return node.ele
}