trpc.group/trpc-go/trpc-go@v1.0.3/naming/loadbalance/random.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 loadbalance
    15  
    16  import (
    17  	"time"
    18  
    19  	"trpc.group/trpc-go/trpc-go/internal/rand"
    20  	"trpc.group/trpc-go/trpc-go/naming/bannednodes"
    21  	"trpc.group/trpc-go/trpc-go/naming/registry"
    22  )
    23  
    24  func init() {
    25  	Register(LoadBalanceRandom, NewRandom())
    26  }
    27  
    28  // Random is the random load balance algorithm.
    29  type Random struct {
    30  	safeRand *rand.SafeRand
    31  }
    32  
    33  // NewRandom creates a new Random.
    34  func NewRandom() *Random {
    35  	return &Random{
    36  		safeRand: rand.NewSafeRand(time.Now().UnixNano()),
    37  	}
    38  }
    39  
    40  // Select picks a node from nodes randomly. Select tries its best to choose a node not in
    41  // bannedNodes of context.
    42  func (b *Random) Select(
    43  	serviceName string,
    44  	nodes []*registry.Node,
    45  	opts ...Option,
    46  ) (node *registry.Node, err error) {
    47  	var o Options
    48  	for _, opt := range opts {
    49  		opt(&o)
    50  	}
    51  
    52  	if o.Ctx == nil {
    53  		return b.chooseOne(nodes)
    54  	}
    55  
    56  	bans, mandatory, ok := bannednodes.FromCtx(o.Ctx)
    57  	if !ok {
    58  		return b.chooseOne(nodes)
    59  	}
    60  
    61  	defer func() {
    62  		if err == nil {
    63  			bannednodes.Add(o.Ctx, node)
    64  		}
    65  	}()
    66  
    67  	node, err = b.chooseUnbanned(nodes, bans)
    68  	if !mandatory && err == ErrNoServerAvailable {
    69  		return b.chooseOne(nodes)
    70  	}
    71  	return node, err
    72  }
    73  
    74  func (b *Random) chooseOne(nodes []*registry.Node) (*registry.Node, error) {
    75  	if len(nodes) == 0 {
    76  		return nil, ErrNoServerAvailable
    77  	}
    78  	return nodes[b.safeRand.Intn(len(nodes))], nil
    79  }
    80  
    81  func (b *Random) chooseUnbanned(
    82  	nodes []*registry.Node,
    83  	bans *bannednodes.Nodes,
    84  ) (*registry.Node, error) {
    85  	if len(nodes) == 0 {
    86  		return nil, ErrNoServerAvailable
    87  	}
    88  	i := b.safeRand.Intn(len(nodes))
    89  	if !bans.Range(func(n *registry.Node) bool {
    90  		return n.Address != nodes[i].Address
    91  	}) {
    92  		return b.chooseUnbanned(append(nodes[:i], nodes[i+1:]...), bans)
    93  	}
    94  	return nodes[i], nil
    95  }