Merge pull request #245 from ydb-platform/234-prefer-local-dc

Detect local DC

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,6 @@
+* Implemented detection of local data-center with measuring tcp dial RTT
+* Added `trace.Driver.OnBalancer{Init,Close,ChooseEndpoint,Update}` events
+* Marked the driver cluster events as deprecated  
 * Simplified the balancing logic
 
 ## v3.25.3

balancers/balancers.go

@@ -24,10 +24,12 @@ func SingleConn() *routerconfig.Config {
 
 // PreferLocalDC creates balancer which use endpoints only in location such as initial endpoint location
 // Balancer "balancer" defines balancing algorithm between endpoints selected with filter by location
+// PreferLocalDC balancer try to autodetect local DC from client side.
 func PreferLocalDC(balancer *routerconfig.Config) *routerconfig.Config {
-	balancer.IsPreferConn = func(c conn.Conn) bool {
-		return c.Endpoint().LocalDC()
+	balancer.IsPreferConn = func(routerInfo routerconfig.Info, c conn.Conn) bool {
+		return c.Endpoint().Location() == routerInfo.SelfLocation
 	}
+	balancer.DetectlocalDC = true
 	return balancer
 }
 
@@ -49,7 +51,7 @@ func PreferLocations(balancer *routerconfig.Config, locations ...string) *router
 	for i := range locations {
 		locations[i] = strings.ToUpper(locations[i])
 	}
-	balancer.IsPreferConn = func(c conn.Conn) bool {
+	balancer.IsPreferConn = func(_ routerconfig.Info, c conn.Conn) bool {
 		location := strings.ToUpper(c.Endpoint().Location())
 		for _, l := range locations {
 			if location == l {
@@ -74,13 +76,16 @@ type Endpoint interface {
 	NodeID() uint32
 	Address() string
 	Location() string
+
+	// Deprecated: LocalDC check "local" by compare endpoint location with discovery "selflocation" field.
+	// It work good only if connection url always point to local dc.
 	LocalDC() bool
 }
 
 // Prefer creates balancer which use endpoints by filter
 // Balancer "balancer" defines balancing algorithm between endpoints selected with filter
 func Prefer(balancer *routerconfig.Config, filter func(endpoint Endpoint) bool) *routerconfig.Config {
-	balancer.IsPreferConn = func(c conn.Conn) bool {
+	balancer.IsPreferConn = func(_ routerconfig.Info, c conn.Conn) bool {
 		return filter(c.Endpoint())
 	}
 	return balancer

balancers/balancers_test.go

@@ -12,24 +12,24 @@ import (
 
 func TestPreferLocalDC(t *testing.T) {
 	conns := []conn.Conn{
-		&mock.Conn{AddrField: "1", LocalDCField: false},
-		&mock.Conn{AddrField: "2", State: conn.Online, LocalDCField: true},
-		&mock.Conn{AddrField: "3", State: conn.Online, LocalDCField: true},
+		&mock.Conn{AddrField: "1", LocationField: "1"},
+		&mock.Conn{AddrField: "2", State: conn.Online, LocationField: "2"},
+		&mock.Conn{AddrField: "3", State: conn.Online, LocationField: "2"},
 	}
-	rr := PreferLocalDC(RoundRobin())
+	rr := PreferLocalDC(Default())
 	require.False(t, rr.AllowFalback)
-	require.Equal(t, []conn.Conn{conns[1], conns[2]}, applyPreferFilter(rr, conns))
+	require.Equal(t, []conn.Conn{conns[1], conns[2]}, applyPreferFilter(routerconfig.Info{SelfLocation: "2"}, rr, conns))
 }
 
 func TestPreferLocalDCWithFallBack(t *testing.T) {
 	conns := []conn.Conn{
-		&mock.Conn{AddrField: "1", LocalDCField: false},
-		&mock.Conn{AddrField: "2", State: conn.Online, LocalDCField: true},
-		&mock.Conn{AddrField: "3", State: conn.Online, LocalDCField: true},
+		&mock.Conn{AddrField: "1", LocationField: "1"},
+		&mock.Conn{AddrField: "2", State: conn.Online, LocationField: "2"},
+		&mock.Conn{AddrField: "3", State: conn.Online, LocationField: "2"},
 	}
-	rr := PreferLocalDCWithFallBack(RoundRobin())
+	rr := PreferLocalDCWithFallBack(Default())
 	require.True(t, rr.AllowFalback)
-	require.Equal(t, []conn.Conn{conns[1], conns[2]}, applyPreferFilter(rr, conns))
+	require.Equal(t, []conn.Conn{conns[1], conns[2]}, applyPreferFilter(routerconfig.Info{SelfLocation: "2"}, rr, conns))
 }
 
 func TestPreferLocations(t *testing.T) {
@@ -39,9 +39,9 @@ func TestPreferLocations(t *testing.T) {
 		&mock.Conn{AddrField: "3", State: conn.Online, LocationField: "two"},
 	}
 
-	rr := PreferLocations(RoundRobin(), "zero", "two")
+	rr := PreferLocations(Default(), "zero", "two")
 	require.False(t, rr.AllowFalback)
-	require.Equal(t, []conn.Conn{conns[0], conns[2]}, applyPreferFilter(rr, conns))
+	require.Equal(t, []conn.Conn{conns[0], conns[2]}, applyPreferFilter(routerconfig.Info{}, rr, conns))
 }
 
 func TestPreferLocationsWithFallback(t *testing.T) {
@@ -51,18 +51,18 @@ func TestPreferLocationsWithFallback(t *testing.T) {
 		&mock.Conn{AddrField: "3", State: conn.Online, LocationField: "two"},
 	}
 
-	rr := PreferLocationsWithFallback(RoundRobin(), "zero", "two")
+	rr := PreferLocationsWithFallback(Default(), "zero", "two")
 	require.True(t, rr.AllowFalback)
-	require.Equal(t, []conn.Conn{conns[0], conns[2]}, applyPreferFilter(rr, conns))
+	require.Equal(t, []conn.Conn{conns[0], conns[2]}, applyPreferFilter(routerconfig.Info{}, rr, conns))
 }
 
-func applyPreferFilter(b *routerconfig.Config, conns []conn.Conn) []conn.Conn {
+func applyPreferFilter(info routerconfig.Info, b *routerconfig.Config, conns []conn.Conn) []conn.Conn {
 	if b.IsPreferConn == nil {
-		b.IsPreferConn = func(c conn.Conn) bool { return true }
+		b.IsPreferConn = func(routerInfo routerconfig.Info, c conn.Conn) bool { return true }
 	}
 	res := make([]conn.Conn, 0, len(conns))
 	for _, c := range conns {
-		if b.IsPreferConn(c) {
+		if b.IsPreferConn(info, c) {
 			res = append(res, c)
 		}
 	}

balancers/config_test.go

@@ -50,10 +50,13 @@ func TestFromConfig(t *testing.T) {
 				"type": "random_choice",
 				"prefer": "local_dc"
 			}`,
-			res: routerconfig.Config{IsPreferConn: func(c conn.Conn) bool {
-				// some non nil func
-				return false
-			}},
+			res: routerconfig.Config{
+				DetectlocalDC: true,
+				IsPreferConn: func(routerInfo routerconfig.Info, c conn.Conn) bool {
+					// some non nil func
+					return false
+				},
+			},
 		},
 		{
 			name: "prefer_unknown_type",
@@ -71,8 +74,9 @@ func TestFromConfig(t *testing.T) {
 				"fallback": true
 			}`,
 			res: routerconfig.Config{
-				AllowFalback: true,
-				IsPreferConn: func(c conn.Conn) bool {
+				AllowFalback:  true,
+				DetectlocalDC: true,
+				IsPreferConn: func(routerInfo routerconfig.Info, c conn.Conn) bool {
 					// some non nil func
 					return false
 				},
@@ -86,7 +90,7 @@ func TestFromConfig(t *testing.T) {
 				"locations": ["AAA", "BBB", "CCC"]
 			}`,
 			res: routerconfig.Config{
-				IsPreferConn: func(c conn.Conn) bool {
+				IsPreferConn: func(routerInfo routerconfig.Info, c conn.Conn) bool {
 					// some non nil func
 					return false
 				},
@@ -102,7 +106,7 @@ func TestFromConfig(t *testing.T) {
 			}`,
 			res: routerconfig.Config{
 				AllowFalback: true,
-				IsPreferConn: func(c conn.Conn) bool {
+				IsPreferConn: func(routerInfo routerconfig.Info, c conn.Conn) bool {
 					// some non nil func
 					return false
 				},

config/config.go

@@ -28,7 +28,7 @@ type Config struct {
 	trace         trace.Driver
 	dialTimeout   time.Duration
 	connectionTTL time.Duration
-	balancer      *routerconfig.Config
+	routerConfig  *routerconfig.Config
 	secure        bool
 	endpoint      string
 	database      string
@@ -106,7 +106,7 @@ func (c Config) Trace() trace.Driver {
 // Balancer is an optional configuration related to selected balancer.
 // That is, some balancing methods allow to be configured.
 func (c Config) Balancer() *routerconfig.Config {
-	return c.balancer
+	return c.routerConfig
 }
 
 // RequestsType set an additional type hint to all requests.
@@ -231,7 +231,7 @@ func WithDialTimeout(timeout time.Duration) Option {
 
 func WithBalancer(balancer *routerconfig.Config) Option {
 	return func(c *Config) {
-		c.balancer = balancer
+		c.routerConfig = balancer
 	}
 }
 
@@ -312,8 +312,8 @@ func defaultConfig() (c Config) {
 		credentials: credentials.NewAnonymousCredentials(
 			credentials.WithSourceInfo("default"),
 		),
-		balancer:  balancers.Default(),
-		tlsConfig: defaultTLSConfig(),
+		routerConfig: balancers.Default(),
+		tlsConfig:    defaultTLSConfig(),
 		grpcOptions: []grpc.DialOption{
 			grpc.WithContextDialer(
 				func(ctx context.Context, address string) (net.Conn, error) {

connection.go

@@ -390,6 +390,17 @@ func open(ctx context.Context, opts ...Option) (_ Connection, err error) {
 		return nil, xerrors.WithStackTrace(errors.New("configuration: empty database"))
 	}
 
+	onDone := trace.DriverOnInit(
+		c.config.Trace(),
+		&ctx,
+		c.config.Endpoint(),
+		c.config.Database(),
+		c.config.Secure(),
+	)
+	defer func() {
+		onDone(err)
+	}()
+
 	if c.pool == nil {
 		c.pool = conn.NewPool(
 			ctx,

internal/ctxlabels/ctxlabels.go

@@ -0,0 +1,16 @@
+package ctxlabels
+
+import "context"
+
+type localDcKey struct{}
+
+func WithLocalDC(ctx context.Context, dc string) context.Context {
+	return context.WithValue(ctx, localDcKey{}, dc)
+}
+
+func ExtractLocalDC(ctx context.Context) string {
+	if val := ctx.Value(localDcKey{}); val != nil {
+		return val.(string)
+	}
+	return ""
+}

internal/router/config/routerconfig.go

@@ -1,15 +1,18 @@
 package routerconfig
 
-// Dedicated package need for prevent cyclo dependencies config -> router -> config
+import "github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
 
-import (
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
-)
+// Dedicated package need for prevent cyclo dependencies config -> router -> config
 
 type Config struct {
-	IsPreferConn PreferConnFunc
-	AllowFalback bool
-	SingleConn   bool
+	IsPreferConn  PreferConnFunc
+	AllowFalback  bool
+	SingleConn    bool
+	DetectlocalDC bool
+}
+
+type Info struct {
+	SelfLocation string
 }
 
-type PreferConnFunc func(c conn.Conn) bool
+type PreferConnFunc func(routerInfo Info, c conn.Conn) bool

internal/router/connections_state.go

@@ -22,6 +22,7 @@ type connectionsState struct {
 func newConnectionsState(
 	conns []conn.Conn,
 	preferFunc routerconfig.PreferConnFunc,
+	routerInfo routerconfig.Info,
 	allowFallback bool,
 ) *connectionsState {
 	res := &connectionsState{
@@ -33,7 +34,7 @@ func newConnectionsState(
 		),
 	}
 
-	res.prefer, res.fallback = sortPreferConnections(conns, preferFunc, allowFallback)
+	res.prefer, res.fallback = sortPreferConnections(conns, preferFunc, routerInfo, allowFallback)
 	if allowFallback {
 		res.all = conns
 	} else {
@@ -130,6 +131,7 @@ func connsToNodeIDMap(conns []conn.Conn) (res map[uint32]conn.Conn) {
 func sortPreferConnections(
 	conns []conn.Conn,
 	preferFunc routerconfig.PreferConnFunc,
+	routerInfo routerconfig.Info,
 	allowFallback bool,
 ) (prefer []conn.Conn, fallback []conn.Conn) {
 	if preferFunc == nil {
@@ -142,7 +144,7 @@ func sortPreferConnections(
 	}
 
 	for _, c := range conns {
-		if preferFunc(c) {
+		if preferFunc(routerInfo, c) {
 			prefer = append(prefer, c)
 		} else if allowFallback {
 			fallback = append(fallback, c)