Merge pull request #257 from ydb-platform/remove-multi-balancer

balanced code has been simplified

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Simplified the balancing logic
+
 ## v3.25.3
 * Changed primary license to `Apache2.0` for auto-detect license 
 * Refactored `types.Struct` value creation

balancers/balancers.go

@@ -4,45 +4,40 @@ import (
 	"strings"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer/multi"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer/rr"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer/single"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
 )
 
+// Deprecated: RoundRobin is RandomChoice now
 func RoundRobin() balancer.Balancer {
-	return rr.RoundRobin(nil)
+	return &balancer.Config{}
 }
 
 func RandomChoice() balancer.Balancer {
-	return rr.RandomChoice(nil)
+	return &balancer.Config{}
 }
 
 func SingleConn() balancer.Balancer {
-	return single.Balancer(nil)
+	return &balancer.Config{
+		SingleConn: true,
+	}
 }
 
 // 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
 func PreferLocalDC(balancer balancer.Balancer) balancer.Balancer {
-	return Prefer(
-		balancer,
-		func(endpoint Endpoint) bool {
-			return endpoint.LocalDC()
-		},
-	)
+	balancer.IsPreferConn = func(c conn.Conn) bool {
+		return c.Endpoint().LocalDC()
+	}
+	return balancer
 }
 
 // PreferLocalDCWithFallBack 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
 // If filter returned zero endpoints from all discovery endpoints list - used all endpoint instead
 func PreferLocalDCWithFallBack(balancer balancer.Balancer) balancer.Balancer {
-	return PreferWithFallback(
-		balancer,
-		func(endpoint Endpoint) bool {
-			return endpoint.LocalDC()
-		},
-	)
+	balancer = PreferLocalDC(balancer)
+	balancer.AllowFalback = true
+	return balancer
 }
 
 // PreferLocations creates balancer which use endpoints only in selected locations (such as "ABC", "DEF", etc.)
@@ -54,42 +49,25 @@ func PreferLocations(balancer balancer.Balancer, locations ...string) balancer.B
 	for i := range locations {
 		locations[i] = strings.ToUpper(locations[i])
 	}
-	return Prefer(
-		balancer,
-		func(endpoint Endpoint) bool {
-			location := strings.ToUpper(endpoint.Location())
-			for _, l := range locations {
-				if location == l {
-					return true
-				}
+	balancer.IsPreferConn = func(c conn.Conn) bool {
+		location := strings.ToUpper(c.Endpoint().Location())
+		for _, l := range locations {
+			if location == l {
+				return true
 			}
-			return false
-		},
-	)
+		}
+		return false
+	}
+	return balancer
 }
 
 // PreferLocationsWithFallback creates balancer which use endpoints only in selected locations
 // Balancer "balancer" defines balancing algorithm between endpoints selected with filter by location
 // If filter returned zero endpoints from all discovery endpoints list - used all endpoint instead
 func PreferLocationsWithFallback(balancer balancer.Balancer, locations ...string) balancer.Balancer {
-	if len(locations) == 0 {
-		panic("empty list of locations")
-	}
-	for i := range locations {
-		locations[i] = strings.ToUpper(locations[i])
-	}
-	return PreferWithFallback(
-		balancer,
-		func(endpoint Endpoint) bool {
-			location := strings.ToUpper(endpoint.Location())
-			for _, l := range locations {
-				if location == l {
-					return true
-				}
-			}
-			return false
-		},
-	)
+	balancer = PreferLocations(balancer, locations...)
+	balancer.AllowFalback = true
+	return balancer
 }
 
 type Endpoint interface {
@@ -102,34 +80,19 @@ type Endpoint interface {
 // Prefer creates balancer which use endpoints by filter
 // Balancer "balancer" defines balancing algorithm between endpoints selected with filter
 func Prefer(balancer balancer.Balancer, filter func(endpoint Endpoint) bool) balancer.Balancer {
-	return multi.Balancer(
-		multi.WithBalancer(
-			balancer,
-			func(cc conn.Conn) bool {
-				return filter(cc.Endpoint())
-			},
-		),
-	)
+	balancer.IsPreferConn = func(c conn.Conn) bool {
+		return filter(c.Endpoint())
+	}
+	return balancer
 }
 
 // PreferWithFallback creates balancer which use endpoints by filter
 // Balancer "balancer" defines balancing algorithm between endpoints selected with filter
 // If filter returned zero endpoints from all discovery endpoints list - used all endpoint instead
 func PreferWithFallback(balancer balancer.Balancer, filter func(endpoint Endpoint) bool) balancer.Balancer {
-	return multi.Balancer(
-		multi.WithBalancer(
-			balancer,
-			func(cc conn.Conn) bool {
-				return filter(cc.Endpoint())
-			},
-		),
-		multi.WithBalancer(
-			balancer,
-			func(cc conn.Conn) bool {
-				return !filter(cc.Endpoint())
-			},
-		),
-	)
+	balancer = Prefer(balancer, filter)
+	balancer.AllowFalback = true
+	return balancer
 }
 
 // Default balancer used by default

balancers/balancers_test.go

@@ -1,94 +1,71 @@
 package balancers
 
 import (
-	"context"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/mock"
+
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer/mock"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
 )
 
 func TestPreferLocalDC(t *testing.T) {
-	ctx := context.Background()
-
 	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},
 	}
-	rr := PreferLocalDC(RoundRobin()).Create(conns)
-	require.Contains(t, []conn.Conn{conns[1], conns[2]}, rr.Next(ctx))
-	require.Contains(t, []conn.Conn{conns[1], conns[2]}, rr.Next(ctx))
-	require.Contains(t, []conn.Conn{conns[1], conns[2]}, rr.Next(ctx))
-
-	// ban local connections
-	conns[1].SetState(conn.Banned)
-	conns[2].SetState(conn.Banned)
-	require.Contains(t, []conn.Conn{conns[1], conns[2]}, rr.Next(ctx, balancer.WithAcceptBanned(true)))
-	require.Nil(t, rr.Next(ctx))
+	rr := PreferLocalDC(RoundRobin())
+	require.False(t, rr.AllowFalback)
+	require.Equal(t, []conn.Conn{conns[1], conns[2]}, applyPreferFilter(rr, conns))
 }
 
 func TestPreferLocalDCWithFallBack(t *testing.T) {
-	ctx := context.Background()
-
 	conns := []conn.Conn{
-		&mock.Conn{AddrField: "1", LocalDCField: false, State: conn.Online},
+		&mock.Conn{AddrField: "1", LocalDCField: false},
 		&mock.Conn{AddrField: "2", State: conn.Online, LocalDCField: true},
 		&mock.Conn{AddrField: "3", State: conn.Online, LocalDCField: true},
 	}
-	rr := PreferLocalDCWithFallBack(RoundRobin()).Create(conns)
-	require.Contains(t, []conn.Conn{conns[1], conns[2]}, rr.Next(ctx))
-	require.Contains(t, []conn.Conn{conns[1], conns[2]}, rr.Next(ctx))
-	require.Contains(t, []conn.Conn{conns[1], conns[2]}, rr.Next(ctx))
-
-	// ban connections
-	conns[1].SetState(conn.Banned)
-	conns[2].SetState(conn.Banned)
-	require.Contains(t, []conn.Conn{conns[1], conns[2]}, rr.Next(ctx, balancer.WithAcceptBanned(true)))
-	require.Equal(t, conns[0], rr.Next(ctx))
+	rr := PreferLocalDCWithFallBack(RoundRobin())
+	require.True(t, rr.AllowFalback)
+	require.Equal(t, []conn.Conn{conns[1], conns[2]}, applyPreferFilter(rr, conns))
 }
 
 func TestPreferLocations(t *testing.T) {
-	ctx := context.Background()
-
 	conns := []conn.Conn{
 		&mock.Conn{AddrField: "1", LocationField: "zero", State: conn.Online},
 		&mock.Conn{AddrField: "2", State: conn.Online, LocationField: "one"},
 		&mock.Conn{AddrField: "3", State: conn.Online, LocationField: "two"},
 	}
 
-	rr := PreferLocations(RoundRobin(), "zero", "two").Create(conns)
-	require.Contains(t, []conn.Conn{conns[0], conns[2]}, rr.Next(ctx))
-	require.Contains(t, []conn.Conn{conns[0], conns[2]}, rr.Next(ctx))
-	require.Contains(t, []conn.Conn{conns[0], conns[2]}, rr.Next(ctx))
-
-	// ban zero, two
-	conns[0].SetState(conn.Banned)
-	conns[2].SetState(conn.Banned)
-	require.Contains(t, []conn.Conn{conns[0], conns[2]}, rr.Next(ctx, balancer.WithAcceptBanned(true)))
-	require.Nil(t, rr.Next(ctx))
+	rr := PreferLocations(RoundRobin(), "zero", "two")
+	require.False(t, rr.AllowFalback)
+	require.Equal(t, []conn.Conn{conns[0], conns[2]}, applyPreferFilter(rr, conns))
 }
 
 func TestPreferLocationsWithFallback(t *testing.T) {
-	ctx := context.Background()
-
 	conns := []conn.Conn{
 		&mock.Conn{AddrField: "1", LocationField: "zero", State: conn.Online},
 		&mock.Conn{AddrField: "2", State: conn.Online, LocationField: "one"},
 		&mock.Conn{AddrField: "3", State: conn.Online, LocationField: "two"},
 	}
 
-	rr := PreferLocationsWithFallback(RoundRobin(), "zero", "two").Create(conns)
-	require.Contains(t, []conn.Conn{conns[0], conns[2]}, rr.Next(ctx))
-	require.Contains(t, []conn.Conn{conns[0], conns[2]}, rr.Next(ctx))
-	require.Contains(t, []conn.Conn{conns[0], conns[2]}, rr.Next(ctx))
+	rr := PreferLocationsWithFallback(RoundRobin(), "zero", "two")
+	require.True(t, rr.AllowFalback)
+	require.Equal(t, []conn.Conn{conns[0], conns[2]}, applyPreferFilter(rr, conns))
+}
 
-	// ban zero, two
-	conns[0].SetState(conn.Banned)
-	conns[2].SetState(conn.Banned)
-	require.Contains(t, []conn.Conn{conns[0], conns[2]}, rr.Next(ctx, balancer.WithAcceptBanned(true)))
-	require.Equal(t, conns[1], rr.Next(ctx))
+func applyPreferFilter(b balancer.Balancer, conns []conn.Conn) []conn.Conn {
+	if b.IsPreferConn == nil {
+		b.IsPreferConn = func(c conn.Conn) bool { return true }
+	}
+	res := make([]conn.Conn, 0, len(conns))
+	for _, c := range conns {
+		if b.IsPreferConn(c) {
+			res = append(res, c)
+		}
+	}
+	return res
 }