trpc.group/trpc-go/trpc-go@v1.0.3/naming/loadbalance/weightroundrobin/weightroundrobin.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  // Package weightroundrobin provides weight round robin utilities.
    15  package weightroundrobin
    16  
    17  import (
    18  	"sync"
    19  	"time"
    20  
    21  	"trpc.group/trpc-go/trpc-go/naming/loadbalance"
    22  	"trpc.group/trpc-go/trpc-go/naming/registry"
    23  )
    24  
    25  var defaultUpdateRate time.Duration = time.Second * 10
    26  
    27  func init() {
    28  	loadbalance.Register("weight_round_robin", NewWeightRoundRobin(defaultUpdateRate))
    29  }
    30  
    31  // NewWeightRoundRobin creates a new WeightRoundRobin.
    32  func NewWeightRoundRobin(interval time.Duration) *WeightRoundRobin {
    33  	if interval == 0 {
    34  		interval = defaultUpdateRate
    35  	}
    36  	return &WeightRoundRobin{
    37  		pickers:  new(sync.Map),
    38  		interval: interval,
    39  	}
    40  }
    41  
    42  // WeightRoundRobin is a smooth weighted roundrobin algorithm.
    43  type WeightRoundRobin struct {
    44  	pickers  *sync.Map
    45  	interval time.Duration
    46  }
    47  
    48  // Select implements loadbalance.LoadBalancer.
    49  func (wrr *WeightRoundRobin) Select(serviceName string, list []*registry.Node,
    50  	opt ...loadbalance.Option) (*registry.Node, error) {
    51  	opts := &loadbalance.Options{}
    52  	for _, o := range opt {
    53  		o(opts)
    54  	}
    55  	p, ok := wrr.pickers.Load(serviceName)
    56  	if ok {
    57  		return p.(*wrrPicker).Pick(list, opts)
    58  	}
    59  
    60  	newPicker := &wrrPicker{
    61  		interval: wrr.interval,
    62  	}
    63  	v, ok := wrr.pickers.LoadOrStore(serviceName, newPicker)
    64  	if !ok {
    65  		return newPicker.Pick(list, opts)
    66  	}
    67  	return v.(*wrrPicker).Pick(list, opts)
    68  }
    69  
    70  // wrrPicker is a picker based on weighted roundrobin algorithm.
    71  type wrrPicker struct {
    72  	list     []*Server
    73  	updated  time.Time
    74  	mu       sync.Mutex
    75  	interval time.Duration
    76  }
    77  
    78  // Server records the node status.
    79  type Server struct {
    80  	node          *registry.Node
    81  	weight        int
    82  	currentWeight int
    83  	effectWeight  int
    84  }
    85  
    86  // Pick picks a node.
    87  func (p *wrrPicker) Pick(list []*registry.Node, opts *loadbalance.Options) (*registry.Node, error) {
    88  	p.mu.Lock()
    89  	defer p.mu.Unlock()
    90  	p.updateState(list)
    91  	if len(p.list) == 0 {
    92  		return nil, loadbalance.ErrNoServerAvailable
    93  	}
    94  	selected := p.selectServer()
    95  	return selected.node, nil
    96  }
    97  
    98  func (p *wrrPicker) selectServer() *Server {
    99  	var selected *Server
   100  	var total int
   101  	for _, s := range p.list {
   102  		s.currentWeight += s.effectWeight
   103  		total += s.effectWeight
   104  		if s.effectWeight < s.weight {
   105  			s.effectWeight++
   106  		}
   107  		if selected == nil || s.currentWeight > selected.currentWeight {
   108  			selected = s
   109  		}
   110  	}
   111  	selected.currentWeight -= total
   112  	return selected
   113  }
   114  
   115  func (p *wrrPicker) updateState(list []*registry.Node) {
   116  	if len(p.list) == 0 ||
   117  		len(p.list) != len(list) ||
   118  		time.Since(p.updated) > p.interval {
   119  		p.list = p.getServers(list)
   120  		p.updated = time.Now()
   121  	}
   122  }
   123  
   124  func (p *wrrPicker) getServers(list []*registry.Node) []*Server {
   125  	servers := make([]*Server, 0, len(list))
   126  	for _, n := range list {
   127  		weight := n.Weight
   128  		if weight == 0 {
   129  			weight = 1000
   130  		}
   131  		servers = append(servers, &Server{
   132  			node:          n,
   133  			weight:        weight,
   134  			effectWeight:  weight,
   135  			currentWeight: 0,
   136  		})
   137  	}
   138  	return servers
   139  }