goconsistencehash/local/hashring.go

package local

import (
	"context"
	"errors"
	"math/rand"
	"time"
)

type HashRing struct {
	Locker
	head *node

	nodes map[string]int
	datas map[string]map[string]struct{}

	rand *rand.Rand
}

type node struct {
	nexts []*node
	score int32
	node  []string
}

func NewHashRing(locker Locker) *HashRing {
	return &HashRing{
		Locker: locker,
		head:   &node{},
		nodes:  make(map[string]int),
		datas:  make(map[string]map[string]struct{}),
		rand:   rand.New(rand.NewSource(time.Now().UnixNano())),
	}
}

func (h *HashRing) search(score int32) *node {
	move := h.head
	for level := len(h.head.nexts) - 1; level >= 0; level-- {
		for move.nexts[level] != nil && move.nexts[level].score < score {
			move = move.nexts[level]
		}

		if move.nexts[level] != nil && move.nexts[level].score == score {
			return move.nexts[level]
		}
	}

	return nil
}

func (h *HashRing) floor(score int32) (int32, bool) {
	if len(h.head.nexts) == 0 {
		return -1, false
	}

	move := h.head
	for level := len(move.nexts) - 1; level >= 0; level-- {
		for move.nexts[level] != nil && move.nexts[level].score < score {
			move = move.nexts[level]
		}
	}

	if move.nexts[0] != nil && move.nexts[0].score == score {
		return move.nexts[0].score, true
	}

	if move == h.head {
		return -1, false
	}

	return move.score, true
}

func (h *HashRing) ceiling(score int32) (int32, bool) {
	if len(h.head.nexts) == 0 {
		return -1, false
	}

	move := h.head
	for level := len(move.nexts) - 1; level >= 0; level-- {
		for move.nexts[level] != nil && move.nexts[level].score < score {
			move = move.nexts[level]
		}
	}

	if move.nexts[0] == nil {
		return -1, false
	}

	return move.nexts[0].score, true
}

func (h *HashRing) first() (int32, bool) {
	if len(h.head.nexts) == 0 {
		return -1, false
	}

	return h.head.nexts[0].score, true
}

func (h *HashRing) last() (int32, bool) {
	if len(h.head.nexts) == 0 {
		return -1, false
	}

	move := h.head
	for level := len(move.nexts) - 1; level >= 0; level-- {
		for move.nexts[level] != nil {
			move = move.nexts[level]
		}
	}

	if move == h.head {
		return -1, false
	}

	return move.score, true
}

func (h *HashRing) roll() int {
	level := 0
	for h.rand.Intn(2) > 0 {
		level++
	}

	return level
}

func (h *HashRing) Add(ctx context.Context, score int32, nodeId string) error {
	nodes := h.search(score)
	if nodes != nil {
		for _, _nodeId := range nodes.node {
			if _nodeId == nodeId {
				return errors.New("alrady inset")
			}
		}

		nodes.node = append(nodes.node, nodeId)
		return nil
	}

	roll := h.roll()
	if len(h.head.nexts)-1 < roll {
		diff := make([]*node, roll+1-len(h.head.nexts))
		h.head.nexts = append(h.head.nexts, diff...)
	}

	newNode := &node{
		score: score,
		node:  []string{nodeId},
		nexts: make([]*node, roll+1),
	}

	move := h.head
	for level := len(move.nexts) - 1; level >= 0; level-- {
		for move.nexts[level] != nil && move.nexts[level].score < score {
			move = move.nexts[level]
		}

		if level > roll {
			continue
		}

		newNode.nexts[level] = move.nexts[level]
		move.nexts[level] = newNode
	}

	return nil
}

func (h *HashRing) Del(ctx context.Context, score int32, nodeId string) error {
	nodes := h.search(score)
	if nodes == nil {
		return errors.New("no such score node")
	}

	index := -1
	for idx, key := range nodes.node {
		if key == nodeId {
			index = idx
			break
		}
	}

	if index == -1 {
		return errors.New("nodes not exists")
	}

	if len(nodes.node) > 1 {
		nodes.node = append(nodes.node[:index], nodes.node[index+1:]...)
		return nil
	}

	move := h.head
	for level := len(move.nexts) - 1; level >= 0; level-- {
		for move.nexts[level] != nil && move.nexts[level].score < score {
			move = move.nexts[level]
		}

		if move.nexts[level] == nil || move.nexts[level].score > score {
			continue
		}

		move.nexts[level] = move.nexts[level].nexts[level]
	}

	diff := 0
	for level := len(h.head.nexts) - 1; level >= 0 && h.head.nexts[level] == nil; level-- {
		diff++
	}

	if diff > 0 {
		h.head.nexts = h.head.nexts[:len(h.head.nexts)-diff]
	}

	return nil
}

func (h *HashRing) Floor(ctx context.Context, score int32) (int32, error) {
	if floorScore, ok := h.floor(score); ok {
		return floorScore, nil
	}

	lastScore, _ := h.last()

	return lastScore, nil
}

func (h *HashRing) Ceiling(ctx context.Context, score int32) (int32, error) {
	if ceillingScore, ok := h.ceiling(score); ok {
		return ceillingScore, nil
	}

	firstScore, _ := h.first()

	return firstScore, nil
}

func (h *HashRing) Node(ctx context.Context, score int32) ([]string, error) {
	node := h.search(score)
	if node == nil {
		return []string{}, errors.New("no such score node")
	}

	return node.node, nil
}

func (h *HashRing) Nodes(context.Context) (map[string]int, error) {
	return h.nodes, nil
}

func (h *HashRing) AddNodes(ctx context.Context, nodeId string, num int) error {
	h.nodes[nodeId] = num

	return nil
}

func (h *HashRing) DelNodes(ctx context.Context, nodeId string) error {
	delete(h.nodes, nodeId)

	return nil
}

func (h *HashRing) Datas(ctx context.Context, nodeId string) (map[string]struct{}, error) {
	return h.datas[nodeId], nil
}

func (h *HashRing) AddDatas(ctx context.Context, nodeId string, datas map[string]struct{}) error {
	datasMap := h.datas[nodeId]
	if datasMap == nil {
		datasMap = make(map[string]struct{})
	}

	for keys := range datas {
		datasMap[keys] = struct{}{}
	}

	h.datas[nodeId] = datasMap

	return nil
}

func (h *HashRing) DelDatas(ctx context.Context, nodeId string, datas map[string]struct{}) error {
	datasMap := h.datas[nodeId]
	if datasMap == nil {
		return nil
	}

	for keys := range datas {
		delete(datasMap, keys)
	}

	if len(datasMap) == 0 {
		delete(h.datas, nodeId)
	}

	return nil
}