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

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