github.com/hdt3213/godis@v1.2.9/lib/consistenthash/consistenthash.go (about)

     1  package consistenthash
     2  
     3  import (
     4  	"hash/crc32"
     5  	"sort"
     6  	"strconv"
     7  	"strings"
     8  )
     9  
    10  // HashFunc defines function to generate hash code
    11  type HashFunc func(data []byte) uint32
    12  
    13  // Map stores nodes and you can pick node from Map
    14  type Map struct {
    15  	hashFunc HashFunc
    16  	replicas int
    17  	keys     []int // sorted
    18  	hashMap  map[int]string
    19  }
    20  
    21  // New creates a new Map
    22  func New(replicas int, fn HashFunc) *Map {
    23  	m := &Map{
    24  		replicas: replicas,
    25  		hashFunc: fn,
    26  		hashMap:  make(map[int]string),
    27  	}
    28  	if m.hashFunc == nil {
    29  		m.hashFunc = crc32.ChecksumIEEE
    30  	}
    31  	return m
    32  }
    33  
    34  // IsEmpty returns if there is no node in Map
    35  func (m *Map) IsEmpty() bool {
    36  	return len(m.keys) == 0
    37  }
    38  
    39  // AddNode add the given nodes into consistent hash circle
    40  func (m *Map) AddNode(keys ...string) {
    41  	for _, key := range keys {
    42  		if key == "" {
    43  			continue
    44  		}
    45  		for i := 0; i < m.replicas; i++ {
    46  			hash := int(m.hashFunc([]byte(strconv.Itoa(i) + key)))
    47  			m.keys = append(m.keys, hash)
    48  			m.hashMap[hash] = key
    49  		}
    50  	}
    51  	sort.Ints(m.keys)
    52  }
    53  
    54  // support hash tag
    55  func getPartitionKey(key string) string {
    56  	beg := strings.Index(key, "{")
    57  	if beg == -1 {
    58  		return key
    59  	}
    60  	end := strings.Index(key, "}")
    61  	if end == -1 || end == beg+1 {
    62  		return key
    63  	}
    64  	return key[beg+1 : end]
    65  }
    66  
    67  // PickNode gets the closest item in the hash to the provided key.
    68  func (m *Map) PickNode(key string) string {
    69  	if m.IsEmpty() {
    70  		return ""
    71  	}
    72  
    73  	partitionKey := getPartitionKey(key)
    74  	hash := int(m.hashFunc([]byte(partitionKey)))
    75  
    76  	// Binary search for appropriate replica.
    77  	idx := sort.Search(len(m.keys), func(i int) bool { return m.keys[i] >= hash })
    78  
    79  	// Means we have cycled back to the first replica.
    80  	if idx == len(m.keys) {
    81  		idx = 0
    82  	}
    83  
    84  	return m.hashMap[m.keys[idx]]
    85  }