trpc.group/trpc-go/trpc-go@v1.0.3/naming/selector/ip_selector.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 selector
    15  
    16  import (
    17  	"errors"
    18  	"strings"
    19  	"time"
    20  
    21  	"trpc.group/trpc-go/trpc-go/internal/rand"
    22  	"trpc.group/trpc-go/trpc-go/naming/bannednodes"
    23  	"trpc.group/trpc-go/trpc-go/naming/discovery"
    24  	"trpc.group/trpc-go/trpc-go/naming/loadbalance"
    25  	"trpc.group/trpc-go/trpc-go/naming/registry"
    26  	"trpc.group/trpc-go/trpc-go/naming/servicerouter"
    27  )
    28  
    29  func init() {
    30  	Register("ip", NewIPSelector())  // ip://ip:port
    31  	Register("dns", NewIPSelector()) // dns://domain:port
    32  }
    33  
    34  // ipSelector is a selector based on ip list.
    35  type ipSelector struct {
    36  	safeRand *rand.SafeRand
    37  }
    38  
    39  // NewIPSelector creates a new ipSelector.
    40  func NewIPSelector() *ipSelector {
    41  	return &ipSelector{
    42  		safeRand: rand.NewSafeRand(time.Now().UnixNano()),
    43  	}
    44  }
    45  
    46  // Select implements Selector.Select. ServiceName may have multiple IP, such as ip1:port1,ip2:port2.
    47  // If ctx has bannedNodes, Select will try its best to select a node not in bannedNodes.
    48  func (s *ipSelector) Select(
    49  	serviceName string, opt ...Option,
    50  ) (node *registry.Node, err error) {
    51  	if serviceName == "" {
    52  		return nil, errors.New("serviceName empty")
    53  	}
    54  
    55  	var o Options = Options{
    56  		DiscoveryOptions:     make([]discovery.Option, 0, defaultDiscoveryOptionsSize),
    57  		ServiceRouterOptions: make([]servicerouter.Option, 0, defaultServiceRouterOptionsSize),
    58  		LoadBalanceOptions:   make([]loadbalance.Option, 0, defaultLoadBalanceOptionsSize),
    59  	}
    60  	for _, opt := range opt {
    61  		opt(&o)
    62  	}
    63  	if o.Ctx == nil {
    64  		addr, err := s.chooseOne(serviceName)
    65  		if err != nil {
    66  			return nil, err
    67  		}
    68  		return &registry.Node{ServiceName: serviceName, Address: addr}, nil
    69  	}
    70  
    71  	bans, mandatory, ok := bannednodes.FromCtx(o.Ctx)
    72  	if !ok {
    73  		addr, err := s.chooseOne(serviceName)
    74  		if err != nil {
    75  			return nil, err
    76  		}
    77  		return &registry.Node{ServiceName: serviceName, Address: addr}, nil
    78  	}
    79  
    80  	defer func() {
    81  		if err == nil {
    82  			bannednodes.Add(o.Ctx, node)
    83  		}
    84  	}()
    85  
    86  	addr, err := s.chooseUnbanned(strings.Split(serviceName, ","), bans)
    87  	if !mandatory && err != nil {
    88  		addr, err = s.chooseOne(serviceName)
    89  	}
    90  	if err != nil {
    91  		return nil, err
    92  	}
    93  	return &registry.Node{ServiceName: serviceName, Address: addr}, nil
    94  }
    95  
    96  func (s *ipSelector) chooseOne(serviceName string) (string, error) {
    97  	num := strings.Count(serviceName, ",") + 1
    98  	if num == 1 {
    99  		return serviceName, nil
   100  	}
   101  
   102  	var addr string
   103  	r := s.safeRand.Intn(num)
   104  	for i := 0; i <= r; i++ {
   105  		j := strings.IndexByte(serviceName, ',')
   106  		if j < 0 {
   107  			addr = serviceName
   108  			break
   109  		}
   110  		addr, serviceName = serviceName[:j], serviceName[j+1:]
   111  	}
   112  	return addr, nil
   113  }
   114  
   115  func (s *ipSelector) chooseUnbanned(addrs []string, bans *bannednodes.Nodes) (string, error) {
   116  	if len(addrs) == 0 {
   117  		return "", errors.New("no available targets")
   118  	}
   119  
   120  	r := s.safeRand.Intn(len(addrs))
   121  	if !bans.Range(func(n *registry.Node) bool {
   122  		return n.Address != addrs[r]
   123  	}) {
   124  		return s.chooseUnbanned(append(addrs[:r], addrs[r+1:]...), bans)
   125  	}
   126  	return addrs[r], nil
   127  }
   128  
   129  // Report reports nothing.
   130  func (s *ipSelector) Report(*registry.Node, time.Duration, error) error {
   131  	return nil
   132  }