trpc.group/trpc-go/trpc-go@v1.0.3/naming/loadbalance/loadbalance.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 is a pluggable loadbalance module.
    15  package loadbalance
    16  
    17  import (
    18  	"errors"
    19  	"sync"
    20  
    21  	"trpc.group/trpc-go/trpc-go/naming/registry"
    22  )
    23  
    24  var (
    25  	// ErrNoServerAvailable is an error that indicate there is no available server.
    26  	ErrNoServerAvailable = errors.New("no server is available")
    27  )
    28  
    29  // load balance strategies.
    30  const (
    31  	LoadBalanceRandom             = "random"
    32  	LoadBalanceRoundRobin         = "round_robin"
    33  	LoadBalanceWeightedRoundRobin = "weight_round_robin"
    34  	LoadBalanceConsistentHash     = "consistent_hash"
    35  )
    36  
    37  // DefaultLoadBalancer is the default LoadBalancer.
    38  var DefaultLoadBalancer LoadBalancer = NewRandom()
    39  
    40  // SetDefaultLoadBalancer sets the default LoadBalancer.
    41  func SetDefaultLoadBalancer(b LoadBalancer) {
    42  	DefaultLoadBalancer = b
    43  }
    44  
    45  // LoadBalancer is the interface that defines the load balance which returns a node from a node list.
    46  type LoadBalancer interface {
    47  	Select(serviceName string, list []*registry.Node, opt ...Option) (node *registry.Node, err error)
    48  }
    49  
    50  var (
    51  	loadbalancers = make(map[string]LoadBalancer)
    52  	lock          = sync.RWMutex{}
    53  )
    54  
    55  // Register registers a named LoadBalancer.
    56  func Register(name string, s LoadBalancer) {
    57  	lock.Lock()
    58  	loadbalancers[name] = s
    59  	lock.Unlock()
    60  }
    61  
    62  // Get gets a named LoadBalancer.
    63  func Get(name string) LoadBalancer {
    64  	lock.RLock()
    65  	lb := loadbalancers[name]
    66  	lock.RUnlock()
    67  	return lb
    68  }