github.com/cloudwego/kitex@v0.9.0/pkg/loadbalance/weighted_round_robin.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  	"context"
    21  	"sync"
    22  
    23  	"github.com/bytedance/gopkg/lang/fastrand"
    24  
    25  	"github.com/cloudwego/kitex/pkg/discovery"
    26  )
    27  
    28  var (
    29  	_ Picker = &WeightedRoundRobinPicker{}
    30  	_ Picker = &RoundRobinPicker{}
    31  )
    32  
    33  const wrrVNodesBatchSize = 500 // it will calculate wrrVNodesBatchSize vnodes when insufficient
    34  
    35  type wrrNode struct {
    36  	discovery.Instance
    37  	current int
    38  }
    39  
    40  func newWeightedRoundRobinPicker(instances []discovery.Instance) Picker {
    41  	wrrp := new(WeightedRoundRobinPicker)
    42  	wrrp.iterator = newRound()
    43  
    44  	// shuffle nodes
    45  	wrrp.size = uint64(len(instances))
    46  	wrrp.nodes = make([]*wrrNode, wrrp.size)
    47  	offset := fastrand.Uint64n(wrrp.size)
    48  	totalWeight := 0
    49  	gcd := 0
    50  	for idx := uint64(0); idx < wrrp.size; idx++ {
    51  		ins := instances[(idx+offset)%wrrp.size]
    52  		totalWeight += ins.Weight()
    53  		gcd = gcdInt(gcd, ins.Weight())
    54  		wrrp.nodes[idx] = &wrrNode{
    55  			Instance: ins,
    56  			current:  0,
    57  		}
    58  	}
    59  
    60  	wrrp.vcapacity = uint64(totalWeight / gcd)
    61  	wrrp.vnodes = make([]discovery.Instance, wrrp.vcapacity)
    62  	wrrp.buildVirtualWrrNodes(wrrVNodesBatchSize)
    63  	return wrrp
    64  }
    65  
    66  // WeightedRoundRobinPicker implement smooth weighted round-robin algorithm.
    67  // Refer from https://github.com/phusion/nginx/commit/27e94984486058d73157038f7950a0a36ecc6e35
    68  type WeightedRoundRobinPicker struct {
    69  	nodes []*wrrNode // instance node with current weight
    70  	size  uint64     // size of wrrNodes
    71  
    72  	iterator  *round
    73  	vsize     uint64               // size of valid vnodes
    74  	vcapacity uint64               // capacity of all vnodes
    75  	vnodes    []discovery.Instance // precalculated vnodes which order by swrr algorithm
    76  	vlock     sync.RWMutex         // mutex for vnodes
    77  }
    78  
    79  // Next implements the Picker interface.
    80  func (wp *WeightedRoundRobinPicker) Next(ctx context.Context, request interface{}) (ins discovery.Instance) {
    81  	// iterator must start from 0, because we need warmup the vnode at beginning
    82  	idx := wp.iterator.Next() % wp.vcapacity
    83  	// fast path
    84  	wp.vlock.RLock()
    85  	ins = wp.vnodes[idx]
    86  	wp.vlock.RUnlock()
    87  	if ins != nil {
    88  		return ins
    89  	}
    90  
    91  	// slow path
    92  	wp.vlock.Lock()
    93  	defer wp.vlock.Unlock()
    94  	if wp.vnodes[idx] != nil { // other goroutine has filled the vnodes
    95  		return wp.vnodes[idx]
    96  	}
    97  	vtarget := wp.vsize + wrrVNodesBatchSize
    98  	// we must promise that idx < vtarget
    99  	if idx >= vtarget {
   100  		vtarget = idx + 1
   101  	}
   102  	wp.buildVirtualWrrNodes(vtarget)
   103  	return wp.vnodes[idx]
   104  }
   105  
   106  func (wp *WeightedRoundRobinPicker) buildVirtualWrrNodes(vtarget uint64) {
   107  	if vtarget > wp.vcapacity {
   108  		vtarget = wp.vcapacity
   109  	}
   110  	for i := wp.vsize; i < vtarget; i++ {
   111  		wp.vnodes[i] = nextWrrNode(wp.nodes).Instance
   112  	}
   113  	wp.vsize = vtarget
   114  }
   115  
   116  func nextWrrNode(nodes []*wrrNode) (selected *wrrNode) {
   117  	maxCurrent := 0
   118  	totalWeight := 0
   119  	for _, node := range nodes {
   120  		node.current += node.Weight()
   121  		totalWeight += node.Weight()
   122  		if selected == nil || node.current > maxCurrent {
   123  			selected = node
   124  			maxCurrent = node.current
   125  		}
   126  	}
   127  	if selected == nil {
   128  		return nil
   129  	}
   130  	selected.current -= totalWeight
   131  	return selected
   132  }
   133  
   134  // RoundRobinPicker .
   135  type RoundRobinPicker struct {
   136  	size      uint64
   137  	instances []discovery.Instance
   138  	iterator  *round
   139  }
   140  
   141  func newRoundRobinPicker(instances []discovery.Instance) Picker {
   142  	size := uint64(len(instances))
   143  	return &RoundRobinPicker{
   144  		size:      size,
   145  		instances: instances,
   146  		iterator:  newRandomRound(),
   147  	}
   148  }
   149  
   150  // Next implements the Picker interface.
   151  func (rp *RoundRobinPicker) Next(ctx context.Context, request interface{}) (ins discovery.Instance) {
   152  	if rp.size == 0 {
   153  		return nil
   154  	}
   155  	idx := rp.iterator.Next() % rp.size
   156  	ins = rp.instances[idx]
   157  	return ins
   158  }
   159  
   160  func gcdInt(a, b int) int {
   161  	for b != 0 {
   162  		a, b = b, a%b
   163  	}
   164  	return a
   165  }