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

     1  package balancer
     2  
     3  import (
     4  	"sort"
     5  	"sync"
     6  
     7  	"github.com/fufuok/balancer/utils"
     8  )
     9  
    10  // WeightedRand
    11  type wr struct {
    12  	items   []*wrItem
    13  	weights []int
    14  	count   int
    15  	max     uint32
    16  	all     map[string]int
    17  
    18  	sync.RWMutex
    19  }
    20  
    21  type wrItem struct {
    22  	item   string
    23  	weight int
    24  }
    25  
    26  func NewWeightedRand(items ...map[string]int) (lb *wr) {
    27  	if len(items) > 0 && len(items[0]) > 0 {
    28  		lb = &wr{}
    29  		lb.Update(items[0])
    30  		return
    31  	}
    32  	return &wr{
    33  		all: make(map[string]int),
    34  	}
    35  }
    36  
    37  func (b *wr) Add(item string, weight ...int) {
    38  	w := 1
    39  	if len(weight) > 0 {
    40  		w = weight[0]
    41  	}
    42  
    43  	all := b.All().(map[string]int)
    44  	all[item] = w
    45  	b.Update(all)
    46  }
    47  
    48  func (b *wr) All() interface{} {
    49  	all := make(map[string]int)
    50  
    51  	b.RLock()
    52  	for k, v := range b.all {
    53  		all[k] = v
    54  	}
    55  	b.RUnlock()
    56  
    57  	return all
    58  }
    59  
    60  func (b *wr) Name() string {
    61  	return "WeightedRand"
    62  }
    63  
    64  func (b *wr) Select(_ ...string) (item string) {
    65  	b.RLock()
    66  	switch b.count {
    67  	case 0:
    68  		item = ""
    69  	case 1:
    70  		item = b.items[0].item
    71  	default:
    72  		r := utils.FastRandn(b.max) + 1
    73  		i := utils.SearchInts(b.weights, int(r))
    74  		item = b.items[i].item
    75  	}
    76  	b.RUnlock()
    77  
    78  	return
    79  }
    80  
    81  func (b *wr) Remove(item string, _ ...bool) (ok bool) {
    82  	b.RLock()
    83  	_, ok = b.all[item]
    84  	b.RUnlock()
    85  
    86  	if ok {
    87  		all := b.All().(map[string]int)
    88  		delete(all, item)
    89  		b.Update(all)
    90  	}
    91  	return
    92  }
    93  
    94  func (b *wr) RemoveAll() {
    95  	b.Lock()
    96  	b.items = b.items[:0]
    97  	b.weights = b.weights[:0]
    98  	b.count = 0
    99  	b.max = 0
   100  	b.all = make(map[string]int)
   101  	b.Unlock()
   102  }
   103  
   104  func (b *wr) Reset() {}
   105  
   106  func (b *wr) Update(items interface{}) bool {
   107  	v, ok := items.(map[string]int)
   108  	if !ok {
   109  		return false
   110  	}
   111  
   112  	var (
   113  		count int
   114  		data  []*wrItem
   115  	)
   116  	for item, weight := range v {
   117  		// Discard items with a weight less than 1
   118  		if weight > 0 {
   119  			data = append(data, &wrItem{
   120  				item:   item,
   121  				weight: weight,
   122  			})
   123  			count++
   124  		}
   125  	}
   126  
   127  	sort.Slice(data, func(i, j int) bool {
   128  		return data[i].weight < data[j].weight
   129  	})
   130  
   131  	max := 0
   132  	weights := make([]int, count)
   133  	for i := range data {
   134  		max += data[i].weight
   135  		weights[i] = max
   136  	}
   137  
   138  	b.Lock()
   139  	b.items = data
   140  	b.weights = weights
   141  	b.count = count
   142  	b.max = uint32(max)
   143  	b.all = v
   144  	b.Unlock()
   145  
   146  	return true
   147  }