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

     1  package balancer
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  // Smooth weighted round-robin balancing
     8  // Ref: https://github.com/phusion/nginx/commit/27e94984486058d73157038f7950a0a36ecc6e35
     9  type swrr struct {
    10  	items []*swrrItem
    11  	count int
    12  	all   map[string]int
    13  
    14  	sync.Mutex
    15  }
    16  
    17  type swrrItem struct {
    18  	item          string
    19  	weight        int
    20  	currentWeight int
    21  }
    22  
    23  func NewSmoothWeightedRoundRobin(items ...map[string]int) (lb *swrr) {
    24  	if len(items) > 0 && len(items[0]) > 0 {
    25  		lb = &swrr{}
    26  		lb.Update(items[0])
    27  		return
    28  	}
    29  	return &swrr{
    30  		all: make(map[string]int),
    31  	}
    32  }
    33  
    34  func (b *swrr) Add(item string, weight ...int) {
    35  	w := 1
    36  	if len(weight) > 0 {
    37  		w = weight[0]
    38  	}
    39  
    40  	b.Lock()
    41  	b.add(item, w)
    42  	b.Unlock()
    43  }
    44  
    45  func (b *swrr) add(item string, weight int) {
    46  	b.remove(item)
    47  	b.items = append(b.items, &swrrItem{
    48  		item:   item,
    49  		weight: weight,
    50  	})
    51  	b.count++
    52  	b.all[item] = weight
    53  }
    54  
    55  func (b *swrr) All() interface{} {
    56  	all := make(map[string]int)
    57  
    58  	b.Lock()
    59  	for k, v := range b.all {
    60  		all[k] = v
    61  	}
    62  	b.Unlock()
    63  
    64  	return all
    65  }
    66  
    67  func (b *swrr) Name() string {
    68  	return "SmoothWeightedRoundRobin"
    69  }
    70  
    71  func (b *swrr) Select(_ ...string) (item string) {
    72  	b.Lock()
    73  	switch b.count {
    74  	case 0:
    75  		item = ""
    76  	case 1:
    77  		if b.items[0].weight > 0 {
    78  			item = b.items[0].item
    79  		}
    80  	default:
    81  		item = b.chooseNext().item
    82  	}
    83  	b.Unlock()
    84  
    85  	return
    86  }
    87  
    88  func (b *swrr) chooseNext() (choice *swrrItem) {
    89  	total := 0
    90  	for i := range b.items {
    91  		c := b.items[i]
    92  		if c == nil {
    93  			return nil
    94  		}
    95  
    96  		total += c.weight
    97  		c.currentWeight += c.weight
    98  
    99  		if choice == nil || c.currentWeight > choice.currentWeight {
   100  			choice = c
   101  		}
   102  	}
   103  
   104  	if choice == nil {
   105  		return nil
   106  	}
   107  
   108  	choice.currentWeight -= total
   109  
   110  	return choice
   111  }
   112  
   113  func (b *swrr) Remove(item string, _ ...bool) bool {
   114  	b.Lock()
   115  	defer b.Unlock()
   116  
   117  	return b.remove(item)
   118  }
   119  
   120  func (b *swrr) remove(item string) (ok bool) {
   121  	for i := 0; i < b.count; i++ {
   122  		if item == b.items[i].item {
   123  			b.items = append(b.items[:i], b.items[i+1:]...)
   124  			b.count--
   125  			delete(b.all, item)
   126  			ok = true
   127  			return
   128  		}
   129  	}
   130  	return
   131  }
   132  
   133  func (b *swrr) RemoveAll() {
   134  	b.Lock()
   135  	b.items = b.items[:0]
   136  	b.count = 0
   137  	b.all = make(map[string]int)
   138  	b.Unlock()
   139  }
   140  
   141  func (b *swrr) Reset() {
   142  	b.Lock()
   143  	for i := range b.items {
   144  		b.items[i].currentWeight = b.items[i].weight
   145  	}
   146  	b.Unlock()
   147  }
   148  
   149  func (b *swrr) Update(items interface{}) bool {
   150  	v, ok := items.(map[string]int)
   151  	if !ok {
   152  		return false
   153  	}
   154  
   155  	count := len(v)
   156  	data := make([]*swrrItem, count)
   157  	i := 0
   158  	for item, weight := range v {
   159  		data[i] = &swrrItem{
   160  			item:   item,
   161  			weight: weight,
   162  		}
   163  		i++
   164  	}
   165  
   166  	b.Lock()
   167  	b.count = count
   168  	b.all = v
   169  	b.items = data
   170  	b.Unlock()
   171  
   172  	return true
   173  }