github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/balancers/balancers.go (about)

     1  package balancers
     2  
     3  import (
     4  	"sort"
     5  	"strings"
     6  
     7  	balancerConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer/config"
     8  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
     9  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xstring"
    10  )
    11  
    12  // Deprecated: RoundRobin is RandomChoice now
    13  func RoundRobin() *balancerConfig.Config {
    14  	return &balancerConfig.Config{}
    15  }
    16  
    17  func RandomChoice() *balancerConfig.Config {
    18  	return &balancerConfig.Config{}
    19  }
    20  
    21  func SingleConn() *balancerConfig.Config {
    22  	return &balancerConfig.Config{
    23  		SingleConn: true,
    24  	}
    25  }
    26  
    27  type filterLocalDC struct{}
    28  
    29  func (filterLocalDC) Allow(info balancerConfig.Info, c conn.Conn) bool {
    30  	return c.Endpoint().Location() == info.SelfLocation
    31  }
    32  
    33  func (filterLocalDC) String() string {
    34  	return "LocalDC"
    35  }
    36  
    37  // PreferLocalDC creates balancer which use endpoints only in location such as initial endpoint location
    38  // Balancer "balancer" defines balancing algorithm between endpoints selected with filter by location
    39  // PreferLocalDC balancer try to autodetect local DC from client side.
    40  func PreferLocalDC(balancer *balancerConfig.Config) *balancerConfig.Config {
    41  	balancer.Filter = filterLocalDC{}
    42  	balancer.DetectLocalDC = true
    43  
    44  	return balancer
    45  }
    46  
    47  // PreferLocalDCWithFallBack creates balancer which use endpoints only in location such as initial endpoint location
    48  // Balancer "balancer" defines balancing algorithm between endpoints selected with filter by location
    49  // If filter returned zero endpoints from all discovery endpoints list - used all endpoint instead
    50  func PreferLocalDCWithFallBack(balancer *balancerConfig.Config) *balancerConfig.Config {
    51  	balancer = PreferLocalDC(balancer)
    52  	balancer.AllowFallback = true
    53  
    54  	return balancer
    55  }
    56  
    57  type filterLocations []string
    58  
    59  func (locations filterLocations) Allow(_ balancerConfig.Info, c conn.Conn) bool {
    60  	location := strings.ToUpper(c.Endpoint().Location())
    61  	for _, l := range locations {
    62  		if location == l {
    63  			return true
    64  		}
    65  	}
    66  
    67  	return false
    68  }
    69  
    70  func (locations filterLocations) String() string {
    71  	buffer := xstring.Buffer()
    72  	defer buffer.Free()
    73  
    74  	buffer.WriteString("Locations{")
    75  	for i, l := range locations {
    76  		if i != 0 {
    77  			buffer.WriteByte(',')
    78  		}
    79  		buffer.WriteString(l)
    80  	}
    81  	buffer.WriteByte('}')
    82  
    83  	return buffer.String()
    84  }
    85  
    86  // PreferLocations creates balancer which use endpoints only in selected locations (such as "ABC", "DEF", etc.)
    87  // Balancer "balancer" defines balancing algorithm between endpoints selected with filter by location
    88  func PreferLocations(balancer *balancerConfig.Config, locations ...string) *balancerConfig.Config {
    89  	if len(locations) == 0 {
    90  		panic("empty list of locations")
    91  	}
    92  	for i := range locations {
    93  		locations[i] = strings.ToUpper(locations[i])
    94  	}
    95  	sort.Strings(locations)
    96  	balancer.Filter = filterLocations(locations)
    97  
    98  	return balancer
    99  }
   100  
   101  // PreferLocationsWithFallback creates balancer which use endpoints only in selected locations
   102  // Balancer "balancer" defines balancing algorithm between endpoints selected with filter by location
   103  // If filter returned zero endpoints from all discovery endpoints list - used all endpoint instead
   104  func PreferLocationsWithFallback(balancer *balancerConfig.Config, locations ...string) *balancerConfig.Config {
   105  	balancer = PreferLocations(balancer, locations...)
   106  	balancer.AllowFallback = true
   107  
   108  	return balancer
   109  }
   110  
   111  type Endpoint interface {
   112  	NodeID() uint32
   113  	Address() string
   114  	Location() string
   115  
   116  	// Deprecated: LocalDC check "local" by compare endpoint location with discovery "selflocation" field.
   117  	// It work good only if connection url always point to local dc.
   118  	LocalDC() bool
   119  }
   120  
   121  type filterFunc func(info balancerConfig.Info, c conn.Conn) bool
   122  
   123  func (p filterFunc) Allow(info balancerConfig.Info, c conn.Conn) bool {
   124  	return p(info, c)
   125  }
   126  
   127  func (p filterFunc) String() string {
   128  	return "Custom"
   129  }
   130  
   131  // Prefer creates balancer which use endpoints by filter
   132  // Balancer "balancer" defines balancing algorithm between endpoints selected with filter
   133  func Prefer(balancer *balancerConfig.Config, filter func(endpoint Endpoint) bool) *balancerConfig.Config {
   134  	balancer.Filter = filterFunc(func(_ balancerConfig.Info, c conn.Conn) bool {
   135  		return filter(c.Endpoint())
   136  	})
   137  
   138  	return balancer
   139  }
   140  
   141  // PreferWithFallback creates balancer which use endpoints by filter
   142  // Balancer "balancer" defines balancing algorithm between endpoints selected with filter
   143  // If filter returned zero endpoints from all discovery endpoints list - used all endpoint instead
   144  func PreferWithFallback(balancer *balancerConfig.Config, filter func(endpoint Endpoint) bool) *balancerConfig.Config {
   145  	balancer = Prefer(balancer, filter)
   146  	balancer.AllowFallback = true
   147  
   148  	return balancer
   149  }
   150  
   151  // Default balancer used by default
   152  func Default() *balancerConfig.Config {
   153  	return RandomChoice()
   154  }