trpc.group/trpc-go/trpc-go@v1.0.3/naming/loadbalance/roundrobin/roundrobin.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 roundrobin provides round robin utilities.
    15  package roundrobin
    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("round_robin", NewRoundRobin(defaultUpdateRate))
    29  }
    30  
    31  // NewRoundRobin creates a new RoundRobin.
    32  func NewRoundRobin(interval time.Duration) *RoundRobin {
    33  	if interval == 0 {
    34  		interval = defaultUpdateRate
    35  	}
    36  	return &RoundRobin{
    37  		pickers:  new(sync.Map),
    38  		interval: interval,
    39  	}
    40  }
    41  
    42  // RoundRobin defines the roundbin.
    43  type RoundRobin struct {
    44  	pickers  *sync.Map
    45  	interval time.Duration
    46  }
    47  
    48  // Select implements loadbalance.LoadBalancer.
    49  func (rr *RoundRobin) 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 := rr.pickers.Load(serviceName)
    56  	if ok {
    57  		return p.(*rrPicker).Pick(list, opts)
    58  	}
    59  
    60  	newPicker := &rrPicker{
    61  		interval: rr.interval,
    62  	}
    63  	v, ok := rr.pickers.LoadOrStore(serviceName, newPicker)
    64  	if !ok {
    65  		return newPicker.Pick(list, opts)
    66  	}
    67  	return v.(*rrPicker).Pick(list, opts)
    68  }
    69  
    70  // rrPicker is a picker based on roundrobin algorithm.
    71  type rrPicker struct {
    72  	list     []*registry.Node
    73  	updated  time.Time
    74  	mu       sync.Mutex
    75  	next     int
    76  	interval time.Duration
    77  }
    78  
    79  // Pick picks a node.
    80  func (p *rrPicker) Pick(list []*registry.Node, opts *loadbalance.Options) (*registry.Node, error) {
    81  	p.mu.Lock()
    82  	defer p.mu.Unlock()
    83  	p.updateState(list)
    84  	if len(p.list) == 0 {
    85  		return nil, loadbalance.ErrNoServerAvailable
    86  	}
    87  	node := p.list[p.next]
    88  	p.next = (p.next + 1) % len(p.list)
    89  	return node, nil
    90  }
    91  
    92  func (p *rrPicker) updateState(list []*registry.Node) {
    93  	if len(p.list) == 0 ||
    94  		len(p.list) != len(list) ||
    95  		time.Since(p.updated) > p.interval {
    96  		p.list = list
    97  		p.updated = time.Now()
    98  		p.next = 0
    99  		return
   100  	}
   101  }