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

     1  package balancer
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/fufuok/balancer/utils"
     7  )
     8  
     9  // Random
    10  type random struct {
    11  	items []string
    12  	count uint32
    13  
    14  	sync.RWMutex
    15  }
    16  
    17  func NewRandom(items ...[]string) (lb *random) {
    18  	lb = &random{}
    19  	if len(items) > 0 && len(items[0]) > 0 {
    20  		lb.Update(items[0])
    21  	}
    22  	return
    23  }
    24  
    25  func (b *random) Add(item string, _ ...int) {
    26  	b.Lock()
    27  	b.items = append(b.items, item)
    28  	b.count++
    29  	b.Unlock()
    30  }
    31  
    32  func (b *random) All() interface{} {
    33  	all := make([]string, b.count)
    34  
    35  	b.Lock()
    36  	for i, v := range b.items {
    37  		all[i] = v
    38  	}
    39  	b.Unlock()
    40  
    41  	return all
    42  }
    43  
    44  func (b *random) Name() string {
    45  	return "Random"
    46  }
    47  
    48  func (b *random) Select(_ ...string) (item string) {
    49  	b.RLock()
    50  	switch b.count {
    51  	case 0:
    52  		item = ""
    53  	case 1:
    54  		item = b.items[0]
    55  	default:
    56  		item = b.items[utils.FastRandn(b.count)]
    57  	}
    58  	b.RUnlock()
    59  
    60  	return
    61  }
    62  
    63  func (b *random) 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 := uint32(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 *random) RemoveAll() {
    84  	b.Lock()
    85  	b.items = b.items[:0]
    86  	b.count = 0
    87  	b.Unlock()
    88  }
    89  
    90  func (b *random) Reset() {}
    91  
    92  func (b *random) Update(items interface{}) bool {
    93  	v, ok := items.([]string)
    94  	if !ok {
    95  		return false
    96  	}
    97  
    98  	b.Lock()
    99  	b.items = v
   100  	b.count = uint32(len(v))
   101  	b.Unlock()
   102  
   103  	return true
   104  }