Merge pull request #191 from ydb-platform/safe-balancers

* Refactored balancers (makes concurrent-safe)

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,11 @@
+* Refactored balancers (makes concurrent-safe)
+* Excluded separate balancers lock from cluster
+* Refactored `cluster.Cluster` interface (`Insert` and `Remove` returning nothing now) 
+* Replaced unsafe `cluster.close` boolean flag to `cluster.done` chan for listening close event
+* Added internal checker `cluster.isClosed()` for check cluster state
+* Extracted getting available conn from balancer to internal helper `cluster.get` (called inside `cluster.Get` as last effort)
+* Added checking `conn.Conn` availability with `conn.Ping()` in prefer nodeID case
+
 ## v3.18.1
 * Added `conn.Ping(ctx)` method for check availability of `conn.Conn`
 * Refactored `cluster.Cluster.Get(ctx)` to return only available connection (instead of returning any connection from balancer)

balancers/balancer.go renamed to balancers/balancers.go

@@ -1,6 +1,8 @@
 package multi
 
 import (
+	"sync"
+
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
 )
@@ -18,6 +20,7 @@ type multiHandle struct {
 }
 
 type multi struct {
+	mu       sync.RWMutex
 	balancer []balancer.Balancer
 	filter   []func(conn.Conn) bool
 }
@@ -35,6 +38,8 @@ func (m *multi) Create() balancer.Balancer {
 
 func WithBalancer(b balancer.Balancer, filter func(cc conn.Conn) bool) Option {
 	return func(m *multi) {
+		m.mu.Lock()
+		defer m.mu.Unlock()
 		m.balancer = append(m.balancer, b)
 		m.filter = append(m.filter, filter)
 	}
@@ -43,6 +48,8 @@ func WithBalancer(b balancer.Balancer, filter func(cc conn.Conn) bool) Option {
 type Option func(*multi)
 
 func (m *multi) Contains(x balancer.Element) bool {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	for i, h := range x.(multiHandle).elements {
 		if h != nil && m.balancer[i].Contains(h) {
 			return true
@@ -52,6 +59,8 @@ func (m *multi) Contains(x balancer.Element) bool {
 }
 
 func (m *multi) Next() conn.Conn {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
 	for _, b := range m.balancer {
 		if c := b.Next(); c != nil {
 			return c
@@ -61,6 +70,9 @@ func (m *multi) Next() conn.Conn {
 }
 
 func (m *multi) Insert(conn conn.Conn) balancer.Element {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	var (
 		n = len(m.filter)
 		h = multiHandle{
@@ -82,6 +94,9 @@ func (m *multi) Insert(conn conn.Conn) balancer.Element {
 }
 
 func (m *multi) Remove(x balancer.Element) (removed bool) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	for i, h := range x.(multiHandle).elements {
 		if h != nil {
 			if m.balancer[i].Remove(h) {

internal/balancer/multi/multi.go

@@ -18,6 +18,7 @@ import (
 // routine – that is, not a runtime metric) and interprets it as inversion of
 // weight.
 type roundRobin struct {
+	mu    sync.RWMutex
 	min   float32
 	max   float32
 	belt  []int
@@ -32,28 +33,29 @@ func (r *roundRobin) Create() balancer.Balancer {
 
 func RoundRobin() balancer.Balancer {
 	return &roundRobin{
-		r: rand.New(),
+		r: rand.New(rand.WithLock()),
 	}
 }
 
 func RandomChoice() balancer.Balancer {
 	return &randomChoice{
 		roundRobin: roundRobin{
-			r: rand.New(),
+			r: rand.New(rand.WithLock()),
 		},
 	}
 }
 
 type randomChoice struct {
 	roundRobin
-	m sync.Mutex
 }
 
 func (r *randomChoice) Create() balancer.Balancer {
 	return RandomChoice()
 }
 
 func (r *roundRobin) Next() conn.Conn {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
 	if n := len(r.conns); n == 0 {
 		return nil
 	}
@@ -63,23 +65,27 @@ func (r *roundRobin) Next() conn.Conn {
 }
 
 func (r *randomChoice) Next() conn.Conn {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
 	if n := len(r.conns); n == 0 {
 		return nil
 	}
-	r.m.Lock()
 	i := r.belt[r.r.Int(len(r.belt))]
-	r.m.Unlock()
 	return r.conns[i].Conn
 }
 
 func (r *roundRobin) Insert(conn conn.Conn) balancer.Element {
+	r.mu.Lock()
+	defer r.mu.Unlock()
 	e := r.conns.Insert(conn)
 	r.updateMinMax(e.Conn)
 	r.belt = r.distribute()
 	return e
 }
 
 func (r *roundRobin) Remove(x balancer.Element) bool {
+	r.mu.Lock()
+	defer r.mu.Unlock()
 	el := x.(*list.Element)
 	r.conns.Remove(el)
 	r.inspectMinMax(el.Conn.Endpoint().LoadFactor())
@@ -88,6 +94,8 @@ func (r *roundRobin) Remove(x balancer.Element) bool {
 }
 
 func (r *roundRobin) Contains(x balancer.Element) bool {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
 	if x == nil {
 		return false
 	}
@@ -98,6 +106,7 @@ func (r *roundRobin) Contains(x balancer.Element) bool {
 	return r.conns.Contains(el)
 }
 
+// r.mu must be held
 func (r *roundRobin) updateMinMax(cc conn.Conn) {
 	if len(r.conns) == 1 {
 		r.min = cc.Endpoint().LoadFactor()
@@ -112,6 +121,7 @@ func (r *roundRobin) updateMinMax(cc conn.Conn) {
 	}
 }
 
+// r.mu must be held
 func (r *roundRobin) inspectMinMax(loadFactor float32) {
 	if r.min != loadFactor && r.max != loadFactor {
 		return
@@ -133,13 +143,15 @@ func (r *roundRobin) inspectMinMax(loadFactor float32) {
 	}
 }
 
+// r.mu must be held
 func (r *roundRobin) distribute() []int {
 	return r.spread(distribution(
 		r.min, int32(len(r.conns)),
 		r.max, 1,
 	))
 }
 
+// r.mu must be held
 func (r *roundRobin) spread(f func(float32) int32) []int {
 	var (
 		dist  = make([]int32, 0, len(r.conns))
@@ -244,13 +256,3 @@ func (h *distItemsHeap) Pop() interface{} {
 	*h = p[:n-1]
 	return x
 }
-
-func IsRoundRobin(i interface{}) bool {
-	_, ok := i.(*roundRobin)
-	return ok
-}
-
-func IsRandomChoice(i interface{}) bool {
-	_, ok := i.(*randomChoice)
-	return ok
-}