github.com/hellofresh/janus@v0.0.0-20230925145208-ce8de8183c67/pkg/proxy/balancer/weight.go (about)

     1  package balancer
     2  
     3  import (
     4  	"errors"
     5  	"math/rand"
     6  )
     7  
     8  type (
     9  	// WeightBalancer balancer
    10  	WeightBalancer struct{}
    11  )
    12  
    13  var (
    14  	// ErrZeroWeight is used when there a zero value weight was given
    15  	ErrZeroWeight = errors.New("invalid backend, weight 0 given")
    16  )
    17  
    18  // NewWeightBalancer creates a new instance of WeightBalancer
    19  func NewWeightBalancer() *WeightBalancer {
    20  	return &WeightBalancer{}
    21  }
    22  
    23  // Elect backend using weight strategy
    24  func (b *WeightBalancer) Elect(hosts []*Target) (*Target, error) {
    25  	if len(hosts) == 0 {
    26  		return nil, ErrEmptyBackendList
    27  	}
    28  
    29  	totalWeight := 0
    30  	for _, host := range hosts {
    31  		totalWeight += host.Weight
    32  	}
    33  
    34  	if totalWeight <= 0 {
    35  		return nil, ErrZeroWeight
    36  	}
    37  
    38  	if len(hosts) == 1 {
    39  		return hosts[0], nil
    40  	}
    41  
    42  	r := rand.Intn(totalWeight)
    43  	pos := 0
    44  
    45  	for _, host := range hosts {
    46  		pos += host.Weight
    47  		if r >= pos {
    48  			continue
    49  		}
    50  		return host, nil
    51  	}
    52  
    53  	return nil, ErrCannotElectBackend
    54  }