github.com/fufuok/balancer@v1.0.0/hash.go (about)

     1  package balancer
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/fufuok/balancer/internal/doublejump"
     7  	"github.com/fufuok/balancer/utils"
     8  )
     9  
    10  // JumpConsistentHash
    11  type consistentHash struct {
    12  	items []string
    13  	count int
    14  	h     *doublejump.Hash
    15  
    16  	sync.RWMutex
    17  }
    18  
    19  func NewConsistentHash(items ...[]string) (lb *consistentHash) {
    20  	if len(items) > 0 && len(items[0]) > 0 {
    21  		lb = &consistentHash{}
    22  		lb.Update(items[0])
    23  		return
    24  	}
    25  	return &consistentHash{
    26  		h: doublejump.NewHash(),
    27  	}
    28  }
    29  
    30  func (b *consistentHash) Add(item string, _ ...int) {
    31  	b.Lock()
    32  	b.items = append(b.items, item)
    33  	b.h.Add(item)
    34  	b.count++
    35  	b.Unlock()
    36  }
    37  
    38  func (b *consistentHash) All() interface{} {
    39  	all := make([]string, b.count)
    40  
    41  	b.Lock()
    42  	for i, v := range b.items {
    43  		all[i] = v
    44  	}
    45  	b.Unlock()
    46  
    47  	return all
    48  }
    49  
    50  func (b *consistentHash) Name() string {
    51  	return "ConsistentHash"
    52  }
    53  
    54  func (b *consistentHash) Select(key ...string) (item string) {
    55  	b.RLock()
    56  	switch b.count {
    57  	case 0:
    58  		item = ""
    59  	case 1:
    60  		item = b.items[0]
    61  	default:
    62  		hash := utils.HashString(key...)
    63  		item, _ = b.h.Get(hash).(string)
    64  	}
    65  	b.RUnlock()
    66  
    67  	return
    68  }
    69  
    70  func (b *consistentHash) Remove(item string, asClean ...bool) (ok bool) {
    71  	b.Lock()
    72  	defer b.Unlock()
    73  
    74  	clean := len(asClean) > 0 && asClean[0]
    75  	for i := 0; i < b.count; i++ {
    76  		if item == b.items[i] {
    77  			b.items = append(b.items[:i], b.items[i+1:]...)
    78  			b.count--
    79  			b.h.Remove(item)
    80  			ok = true
    81  			// remove all or remove one
    82  			if !clean {
    83  				return
    84  			}
    85  			i--
    86  		}
    87  	}
    88  	return
    89  }
    90  
    91  func (b *consistentHash) RemoveAll() {
    92  	b.Lock()
    93  	b.items = b.items[:0]
    94  	b.count = 0
    95  	b.h = doublejump.NewHash()
    96  	b.Unlock()
    97  }
    98  
    99  func (b *consistentHash) Reset() {}
   100  
   101  func (b *consistentHash) Update(items interface{}) bool {
   102  	v, ok := items.([]string)
   103  	if !ok {
   104  		return false
   105  	}
   106  
   107  	h := doublejump.NewHash()
   108  	for _, x := range v {
   109  		h.Add(x)
   110  	}
   111  
   112  	b.Lock()
   113  	b.count = len(v)
   114  	b.items = v
   115  	b.h = h
   116  	b.Unlock()
   117  
   118  	return true
   119  }