github.com/cloudwego/kitex@v0.9.0/pkg/loadbalance/weighted_balancer.go (about)

     1  /*
     2   * Copyright 2023 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package loadbalance
    18  
    19  import (
    20  	"sync"
    21  
    22  	"golang.org/x/sync/singleflight"
    23  
    24  	"github.com/cloudwego/kitex/pkg/discovery"
    25  	"github.com/cloudwego/kitex/pkg/klog"
    26  )
    27  
    28  const (
    29  	lbKindRoundRobin = iota
    30  	lbKindInterleaved
    31  	lbKindRandom
    32  	lbKindRandomWithAliasMethod
    33  )
    34  
    35  type weightedBalancer struct {
    36  	kind        int
    37  	pickerCache sync.Map
    38  	sfg         singleflight.Group
    39  }
    40  
    41  // NewWeightedBalancer creates a loadbalancer using weighted-round-robin algorithm.
    42  func NewWeightedBalancer() Loadbalancer {
    43  	return NewWeightedRoundRobinBalancer()
    44  }
    45  
    46  // NewWeightedRoundRobinBalancer creates a loadbalancer using weighted-round-robin algorithm.
    47  func NewWeightedRoundRobinBalancer() Loadbalancer {
    48  	lb := &weightedBalancer{kind: lbKindRoundRobin}
    49  	return lb
    50  }
    51  
    52  // NewInterleavedWeightedRoundRobinBalancer creates a loadbalancer using interleaved-weighted-round-robin algorithm.
    53  func NewInterleavedWeightedRoundRobinBalancer() Loadbalancer {
    54  	lb := &weightedBalancer{kind: lbKindInterleaved}
    55  	return lb
    56  }
    57  
    58  // NewWeightedRandomBalancer creates a loadbalancer using weighted-random algorithm.
    59  func NewWeightedRandomBalancer() Loadbalancer {
    60  	lb := &weightedBalancer{kind: lbKindRandom}
    61  	return lb
    62  }
    63  
    64  // NewWeightedRandomWithAliasMethodBalancer creates a loadbalancer using alias-method algorithm.
    65  func NewWeightedRandomWithAliasMethodBalancer() Loadbalancer {
    66  	lb := &weightedBalancer{kind: lbKindRandomWithAliasMethod}
    67  	return lb
    68  }
    69  
    70  // GetPicker implements the Loadbalancer interface.
    71  func (wb *weightedBalancer) GetPicker(e discovery.Result) Picker {
    72  	if !e.Cacheable {
    73  		picker := wb.createPicker(e)
    74  		return picker
    75  	}
    76  
    77  	picker, ok := wb.pickerCache.Load(e.CacheKey)
    78  	if !ok {
    79  		picker, _, _ = wb.sfg.Do(e.CacheKey, func() (interface{}, error) {
    80  			p := wb.createPicker(e)
    81  			wb.pickerCache.Store(e.CacheKey, p)
    82  			return p, nil
    83  		})
    84  	}
    85  	return picker.(Picker)
    86  }
    87  
    88  func (wb *weightedBalancer) createPicker(e discovery.Result) (picker Picker) {
    89  	instances := make([]discovery.Instance, len(e.Instances)) // removed zero weight instances
    90  	weightSum := 0
    91  	balance := true
    92  	cnt := 0
    93  	for idx, instance := range e.Instances {
    94  		weight := instance.Weight()
    95  		if weight <= 0 {
    96  			klog.Warnf("KITEX: invalid weight, weight=%d instance=%s", weight, e.Instances[idx].Address())
    97  			continue
    98  		}
    99  		weightSum += weight
   100  		instances[cnt] = instance
   101  		if cnt > 0 && instances[cnt-1].Weight() != weight {
   102  			balance = false
   103  		}
   104  		cnt++
   105  	}
   106  	instances = instances[:cnt]
   107  	if len(instances) == 0 {
   108  		return new(DummyPicker)
   109  	}
   110  
   111  	switch wb.kind {
   112  	case lbKindRoundRobin:
   113  		if balance {
   114  			picker = newRoundRobinPicker(instances)
   115  		} else {
   116  			picker = newWeightedRoundRobinPicker(instances)
   117  		}
   118  	case lbKindInterleaved:
   119  		if balance {
   120  			picker = newRoundRobinPicker(instances)
   121  		} else {
   122  			picker = newInterleavedWeightedRoundRobinPicker(instances)
   123  		}
   124  	case lbKindRandomWithAliasMethod:
   125  		if balance {
   126  			picker = newRandomPicker(instances)
   127  		} else {
   128  			picker = newAliasMethodPicker(instances, weightSum)
   129  		}
   130  	default: // random
   131  		if balance {
   132  			picker = newRandomPicker(instances)
   133  		} else {
   134  			picker = newWeightedRandomPickerWithSum(instances, weightSum)
   135  		}
   136  	}
   137  	return picker
   138  }
   139  
   140  // Rebalance implements the Rebalancer interface.
   141  func (wb *weightedBalancer) Rebalance(change discovery.Change) {
   142  	if !change.Result.Cacheable {
   143  		return
   144  	}
   145  	wb.pickerCache.Store(change.Result.CacheKey, wb.createPicker(change.Result))
   146  }
   147  
   148  // Delete implements the Rebalancer interface.
   149  func (wb *weightedBalancer) Delete(change discovery.Change) {
   150  	if !change.Result.Cacheable {
   151  		return
   152  	}
   153  	wb.pickerCache.Delete(change.Result.CacheKey)
   154  }
   155  
   156  func (wb *weightedBalancer) Name() string {
   157  	switch wb.kind {
   158  	case lbKindRoundRobin:
   159  		return "weight_round_robin"
   160  	case lbKindInterleaved:
   161  		return "interleaved_weighted_round_robin"
   162  	case lbKindRandomWithAliasMethod:
   163  		return "weight_random_with_alias_method"
   164  	default:
   165  		return "weight_random"
   166  	}
   167  }