11package balancers
22
33import (
4+ "sort"
45 "strings"
56
67 balancerConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer/config"
78 "github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
9+ "github.com/ydb-platform/ydb-go-sdk/v3/internal/xstring"
810)
911
1012// Deprecated: RoundRobin is RandomChoice now
@@ -22,14 +24,22 @@ func SingleConn() *balancerConfig.Config {
2224 }
2325}
2426
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+
2537// PreferLocalDC creates balancer which use endpoints only in location such as initial endpoint location
2638// Balancer "balancer" defines balancing algorithm between endpoints selected with filter by location
2739// PreferLocalDC balancer try to autodetect local DC from client side.
2840func PreferLocalDC (balancer * balancerConfig.Config ) * balancerConfig.Config {
29- balancer .IsPreferConn = func (info balancerConfig.Info , c conn.Conn ) bool {
30- return c .Endpoint ().Location () == info .SelfLocation
31- }
32- balancer .DetectlocalDC = true
41+ balancer .Filter = filterLocalDC {}
42+ balancer .DetectLocalDC = true
3343 return balancer
3444}
3545
@@ -38,10 +48,38 @@ func PreferLocalDC(balancer *balancerConfig.Config) *balancerConfig.Config {
3848// If filter returned zero endpoints from all discovery endpoints list - used all endpoint instead
3949func PreferLocalDCWithFallBack (balancer * balancerConfig.Config ) * balancerConfig.Config {
4050 balancer = PreferLocalDC (balancer )
41- balancer .AllowFalback = true
51+ balancer .AllowFallback = true
4252 return balancer
4353}
4454
55+ type filterLocations []string
56+
57+ func (locations filterLocations ) Allow (_ balancerConfig.Info , c conn.Conn ) bool {
58+ location := strings .ToUpper (c .Endpoint ().Location ())
59+ for _ , l := range locations {
60+ if location == l {
61+ return true
62+ }
63+ }
64+ return false
65+ }
66+
67+ func (locations filterLocations ) String () string {
68+ buffer := xstring .Buffer ()
69+ defer buffer .Free ()
70+
71+ buffer .WriteString ("Locations{" )
72+ for i , l := range locations {
73+ if i != 0 {
74+ buffer .WriteByte (',' )
75+ }
76+ buffer .WriteString (l )
77+ }
78+ buffer .WriteByte ('}' )
79+
80+ return buffer .String ()
81+ }
82+
4583// PreferLocations creates balancer which use endpoints only in selected locations (such as "ABC", "DEF", etc.)
4684// Balancer "balancer" defines balancing algorithm between endpoints selected with filter by location
4785func PreferLocations (balancer * balancerConfig.Config , locations ... string ) * balancerConfig.Config {
@@ -51,15 +89,8 @@ func PreferLocations(balancer *balancerConfig.Config, locations ...string) *bala
5189 for i := range locations {
5290 locations [i ] = strings .ToUpper (locations [i ])
5391 }
54- balancer .IsPreferConn = func (_ balancerConfig.Info , c conn.Conn ) bool {
55- location := strings .ToUpper (c .Endpoint ().Location ())
56- for _ , l := range locations {
57- if location == l {
58- return true
59- }
60- }
61- return false
62- }
92+ sort .Strings (locations )
93+ balancer .Filter = filterLocations (locations )
6394 return balancer
6495}
6596
@@ -68,7 +99,7 @@ func PreferLocations(balancer *balancerConfig.Config, locations ...string) *bala
6899// If filter returned zero endpoints from all discovery endpoints list - used all endpoint instead
69100func PreferLocationsWithFallback (balancer * balancerConfig.Config , locations ... string ) * balancerConfig.Config {
70101 balancer = PreferLocations (balancer , locations ... )
71- balancer .AllowFalback = true
102+ balancer .AllowFallback = true
72103 return balancer
73104}
74105
@@ -82,12 +113,22 @@ type Endpoint interface {
82113 LocalDC () bool
83114}
84115
116+ type filterFunc func (info balancerConfig.Info , c conn.Conn ) bool
117+
118+ func (p filterFunc ) Allow (info balancerConfig.Info , c conn.Conn ) bool {
119+ return p (info , c )
120+ }
121+
122+ func (p filterFunc ) String () string {
123+ return "Custom"
124+ }
125+
85126// Prefer creates balancer which use endpoints by filter
86127// Balancer "balancer" defines balancing algorithm between endpoints selected with filter
87128func Prefer (balancer * balancerConfig.Config , filter func (endpoint Endpoint ) bool ) * balancerConfig.Config {
88- balancer .IsPreferConn = func (_ balancerConfig.Info , c conn.Conn ) bool {
129+ balancer .Filter = filterFunc ( func (_ balancerConfig.Info , c conn.Conn ) bool {
89130 return filter (c .Endpoint ())
90- }
131+ })
91132 return balancer
92133}
93134
@@ -96,7 +137,7 @@ func Prefer(balancer *balancerConfig.Config, filter func(endpoint Endpoint) bool
96137// If filter returned zero endpoints from all discovery endpoints list - used all endpoint instead
97138func PreferWithFallback (balancer * balancerConfig.Config , filter func (endpoint Endpoint ) bool ) * balancerConfig.Config {
98139 balancer = Prefer (balancer , filter )
99- balancer .AllowFalback = true
140+ balancer .AllowFallback = true
100141 return balancer
101142}
102143
0 commit comments