balancers.FromConfig

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,6 @@
+## 3.14.0
+* Added `balacers.FromConfig` balancers creator
+
 ## 3.13.3
 * Fixed linter issues
 

balancers/config.go

@@ -0,0 +1,104 @@
+package balancers
+
+import (
+	"encoding/json"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/errors"
+)
+
+type balancerType string
+
+const (
+	typeRoundRobin   = balancerType("round_robin")
+	typeRandomChoice = balancerType("random_choice")
+	typeSingle       = balancerType("single")
+)
+
+type preferType string
+
+const (
+	preferLocalDC   = preferType("local_dc")
+	preferLocations = preferType("locations")
+)
+
+type balancersConfig struct {
+	Type      balancerType `json:"type"`
+	Prefer    preferType   `json:"prefer,omitempty"`
+	Fallback  bool         `json:"fallback,omitempty"`
+	Locations []string     `json:"locations,omitempty"`
+}
+
+type fromConfigOptionsHolder struct {
+	fallbackBalancer balancer.Balancer
+	errorHandler     func(error)
+}
+
+type fromConfigOption func(h *fromConfigOptionsHolder)
+
+func WithParseErrorFallbackBalancer(b balancer.Balancer) fromConfigOption {
+	return func(h *fromConfigOptionsHolder) {
+		h.fallbackBalancer = b
+	}
+}
+
+func WithParseErrorHandler(errorHandler func(error)) fromConfigOption {
+	return func(h *fromConfigOptionsHolder) {
+		h.errorHandler = errorHandler
+	}
+}
+
+func FromConfig(config string, opts ...fromConfigOption) (b balancer.Balancer) {
+	var (
+		h = fromConfigOptionsHolder{
+			fallbackBalancer: Default(),
+		}
+		c balancersConfig
+	)
+
+	for _, o := range opts {
+		o(&h)
+	}
+
+	if err := json.Unmarshal([]byte(config), &c); err != nil {
+		if h.errorHandler != nil {
+			h.errorHandler(err)
+		}
+		return h.fallbackBalancer
+	}
+
+	switch c.Type {
+	case typeSingle:
+		b = SingleConn()
+	case typeRandomChoice:
+		b = RandomChoice()
+	case typeRoundRobin:
+		b = RoundRobin()
+	default:
+		if h.errorHandler != nil {
+			h.errorHandler(errors.Errorf("unknown type of balancer: %s", c.Type))
+		}
+		return h.fallbackBalancer
+	}
+
+	switch c.Prefer {
+	case preferLocalDC:
+		if c.Fallback {
+			return PreferLocalDCWithFallBack(b)
+		}
+		return PreferLocalDC(b)
+	case preferLocations:
+		if len(c.Locations) == 0 {
+			if h.errorHandler != nil {
+				h.errorHandler(errors.Errorf("empty locations list in balancer '%s' config", c.Type))
+			}
+			return h.fallbackBalancer
+		}
+		if c.Fallback {
+			return PreferLocationsWithFallback(b, c.Locations...)
+		}
+		return PreferLocations(b, c.Locations...)
+	default:
+		return b
+	}
+}

balancers/config_test.go

@@ -0,0 +1,126 @@
+package balancers
+
+import (
+	"testing"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint/info"
+)
+
+type testBalancer struct{}
+
+func (t testBalancer) Next() conn.Conn {
+	panic("unexpected call")
+}
+
+func (t testBalancer) Insert(conn conn.Conn) balancer.Element {
+	panic("unexpected call")
+}
+
+func (t testBalancer) Update(element balancer.Element, info info.Info) {
+	panic("unexpected call")
+}
+
+func (t testBalancer) Remove(element balancer.Element) bool {
+	panic("unexpected call")
+}
+
+func (t testBalancer) Contains(element balancer.Element) bool {
+	panic("unexpected call")
+}
+
+func TestFromConfig(t *testing.T) {
+	for _, test := range []struct {
+		name   string
+		config string
+		fail   bool
+	}{
+		{
+			name:   "empty",
+			config: ``,
+			fail:   true,
+		},
+		{
+			name: "single",
+			config: `{
+				"type": "single"
+			}`,
+		},
+		{
+			name: "round_robin",
+			config: `{
+				"type": "round_robin"
+			}`,
+		},
+		{
+			name: "random_choice",
+			config: `{
+				"type": "random_choice"
+			}`,
+		},
+		{
+			name: "prefer_local_dc",
+			config: `{
+				"type": "random_choice",
+				"prefer": "local_dc"
+			}`,
+		},
+		{
+			name: "prefer_unknown_type",
+			config: `{
+				"type": "unknown_type",
+				"prefer": "local_dc"
+			}`,
+			fail: true,
+		},
+		{
+			name: "prefer_local_dc_with_fallback",
+			config: `{
+				"type": "random_choice",
+				"prefer": "local_dc",
+				"fallback": true
+			}`,
+		},
+		{
+			name: "prefer_locations",
+			config: `{
+				"type": "random_choice",
+				"prefer": "locations",
+				"locations": ["AAA", "BBB", "CCC"]
+			}`,
+		},
+		{
+			name: "prefer_locations_with_fallback",
+			config: `{
+				"type": "random_choice",
+				"prefer": "locations",
+				"locations": ["AAA", "BBB", "CCC"],
+				"fallback": true
+			}`,
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			var (
+				actErr   error
+				fallback testBalancer
+			)
+			b := FromConfig(
+				test.config,
+				WithParseErrorFallbackBalancer(fallback),
+				WithParseErrorHandler(func(err error) {
+					actErr = err
+				}),
+			)
+			if test.fail && actErr == nil {
+				t.Fatalf("expected error, but it not hanled")
+			}
+			if !test.fail && actErr != nil {
+				t.Fatalf("unexpected error: %v", actErr)
+			}
+			if test.fail && b != fallback {
+				t.Fatalf("unexpected balancer: %v", b)
+			}
+		})
+	}
+}

internal/meta/version.go

@@ -1,5 +1,5 @@
 package meta
 
 const (
-	Version = "ydb-go-sdk/3.13.3"
+	Version = "ydb-go-sdk/3.14.0"
 )