lock cluster for update endpoints

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,6 @@
+## 3.11.4
+* Refactored `internal/cluster.Cluster` (add option for notify about external lock, lock cluster for update cluster endpoints)
+
 ## 3.11.3
 * Replaced in `table/types/compare_test.go` checking error by error message to checking with `errors.Is()`
 

internal/cluster/cluster.go

@@ -57,6 +57,14 @@ type cluster struct {
 	closed bool
 }
 
+func (c *cluster) Lock() {
+	c.mu.Lock()
+}
+
+func (c *cluster) Unlock() {
+	c.mu.Unlock()
+}
+
 func (c *cluster) GetConn(endpoint endpoint.Endpoint) conn.Conn {
 	return c.pool.GetConn(endpoint)
 }
@@ -69,11 +77,38 @@ func (c *cluster) SetExplorer(repeater repeater.Repeater) {
 	c.explorer = repeater
 }
 
+type crudOptionsHolder struct {
+	locked bool
+}
+
+type crudOption func(h *crudOptionsHolder)
+
+func WithoutLock() crudOption {
+	return func(h *crudOptionsHolder) {
+		h.locked = true
+	}
+}
+
+func parseOptions(opts ...crudOption) *crudOptionsHolder {
+	h := &crudOptionsHolder{}
+	for _, o := range opts {
+		o(h)
+	}
+	return h
+}
+
 type CRUD interface {
-	Insert(ctx context.Context, endpoint endpoint.Endpoint) conn.Conn
-	Update(ctx context.Context, endpoint endpoint.Endpoint) conn.Conn
-	Remove(ctx context.Context, endpoint endpoint.Endpoint) conn.Conn
-	Get(ctx context.Context) (cc conn.Conn, err error)
+	// Insert inserts endpoint to cluster
+	Insert(ctx context.Context, endpoint endpoint.Endpoint, opts ...crudOption) conn.Conn
+
+	// Update updates endpoint in cluster
+	Update(ctx context.Context, endpoint endpoint.Endpoint, opts ...crudOption) conn.Conn
+
+	// Remove removes endpoint from cluster
+	Remove(ctx context.Context, endpoint endpoint.Endpoint, opts ...crudOption) conn.Conn
+
+	// Get gets conn from cluster
+	Get(ctx context.Context, opts ...crudOption) (cc conn.Conn, err error)
 }
 
 type Pessimizer interface {
@@ -85,16 +120,23 @@ type Explorer interface {
 	Force()
 }
 
-type CRUDExplorer interface {
+type Locker interface {
+	Lock()
+	Unlock()
+}
+
+type CRUDExplorerLocker interface {
 	CRUD
 	Explorer
+	Locker
 }
 
 type Cluster interface {
 	closer.Closer
 	CRUD
 	Pessimizer
 	Explorer
+	Locker
 	conn.PoolGetter
 }
 
@@ -151,7 +193,7 @@ func (c *cluster) Close(ctx context.Context) (err error) {
 
 // Get returns next available connection.
 // It returns error on given deadline cancellation or when cluster become closed.
-func (c *cluster) Get(ctx context.Context) (cc conn.Conn, err error) {
+func (c *cluster) Get(ctx context.Context, opts ...crudOption) (cc conn.Conn, err error) {
 	var cancel context.CancelFunc
 	ctx, cancel = context.WithTimeout(ctx, MaxGetConnTimeout)
 	defer cancel()
@@ -192,7 +234,7 @@ func (c *cluster) Get(ctx context.Context) (cc conn.Conn, err error) {
 }
 
 // Insert inserts new connection into the cluster.
-func (c *cluster) Insert(ctx context.Context, e endpoint.Endpoint) (cc conn.Conn) {
+func (c *cluster) Insert(ctx context.Context, e endpoint.Endpoint, opts ...crudOption) (cc conn.Conn) {
 	onDone := trace.DriverOnClusterInsert(c.config.Trace(), &ctx, e)
 	defer func() {
 		if cc != nil {
@@ -202,8 +244,11 @@ func (c *cluster) Insert(ctx context.Context, e endpoint.Endpoint) (cc conn.Conn
 		}
 	}()
 
-	c.mu.Lock()
-	defer c.mu.Unlock()
+	options := parseOptions(opts...)
+	if !options.locked {
+		c.mu.Lock()
+		defer c.mu.Unlock()
+	}
 
 	if c.closed {
 		return nil
@@ -234,7 +279,7 @@ func (c *cluster) Insert(ctx context.Context, e endpoint.Endpoint) (cc conn.Conn
 }
 
 // Update updates existing connection's runtime stats such that load factor and others.
-func (c *cluster) Update(ctx context.Context, e endpoint.Endpoint) (cc conn.Conn) {
+func (c *cluster) Update(ctx context.Context, e endpoint.Endpoint, opts ...crudOption) (cc conn.Conn) {
 	onDone := trace.DriverOnClusterUpdate(c.config.Trace(), &ctx, e)
 	defer func() {
 		if cc != nil {
@@ -244,8 +289,12 @@ func (c *cluster) Update(ctx context.Context, e endpoint.Endpoint) (cc conn.Conn
 		}
 	}()
 
-	c.mu.Lock()
-	defer c.mu.Unlock()
+	options := parseOptions(opts...)
+	if !options.locked {
+		c.mu.Lock()
+		defer c.mu.Unlock()
+	}
+
 	if c.closed {
 		return
 	}
@@ -260,9 +309,11 @@ func (c *cluster) Update(ctx context.Context, e endpoint.Endpoint) (cc conn.Conn
 
 	delete(c.endpoints, e.NodeID())
 	c.index[e.Address()] = entry
+
 	if e.NodeID() > 0 {
 		c.endpoints[e.NodeID()] = entry.Conn
 	}
+
 	if entry.Handle != nil {
 		// entry.Handle may be nil when connection is being tracked.
 		c.balancer.Update(entry.Handle, info.Info{})
@@ -272,7 +323,7 @@ func (c *cluster) Update(ctx context.Context, e endpoint.Endpoint) (cc conn.Conn
 }
 
 // Remove removes and closes previously inserted connection.
-func (c *cluster) Remove(ctx context.Context, e endpoint.Endpoint) (cc conn.Conn) {
+func (c *cluster) Remove(ctx context.Context, e endpoint.Endpoint, opts ...crudOption) (cc conn.Conn) {
 	onDone := trace.DriverOnClusterRemove(c.config.Trace(), &ctx, e)
 	defer func() {
 		if cc != nil {
@@ -282,9 +333,15 @@ func (c *cluster) Remove(ctx context.Context, e endpoint.Endpoint) (cc conn.Conn
 		}
 	}()
 
-	c.mu.Lock()
+	options := parseOptions(opts...)
+	if !options.locked {
+		c.mu.Lock()
+	}
+
 	if c.closed {
-		c.mu.Unlock()
+		if !options.locked {
+			c.mu.Unlock()
+		}
 		return
 	}
 
@@ -298,7 +355,9 @@ func (c *cluster) Remove(ctx context.Context, e endpoint.Endpoint) (cc conn.Conn
 	delete(c.index, e.Address())
 	delete(c.endpoints, e.NodeID())
 
-	c.mu.Unlock()
+	if !options.locked {
+		c.mu.Unlock()
+	}
 
 	if entry.Conn != nil {
 		// entry.Conn may be nil when connection is being tracked after unsuccessful dial().

internal/conn/pool.go

@@ -68,19 +68,23 @@ func (p *pool) GetConn(e endpoint.Endpoint) Conn {
 
 func (p *pool) Close(ctx context.Context) error {
 	close(p.done)
-	var issues []error
+
 	p.mtx.Lock()
 	defer p.mtx.Unlock()
+
+	var issues []error
 	for a, c := range p.conns {
 		if err := c.Close(ctx); err != nil {
 			issues = append(issues, err)
 		}
 		delete(p.conns, a)
 	}
-	if len(issues) == 0 {
-		return nil
+
+	if len(issues) > 0 {
+		return errors.NewWithIssues("connection pool close failed", issues...)
 	}
-	return errors.NewWithIssues("connection pool close failed", issues...)
+
+	return nil
 }
 
 func (p *pool) connCloser(ctx context.Context, interval time.Duration) {

internal/discovery/discovery.go

@@ -23,16 +23,19 @@ import (
 func New(
 	ctx context.Context,
 	cc conn.Conn,
-	crudExplorer cluster.CRUDExplorer,
+	crudExplorer cluster.CRUDExplorerLocker,
 	opts ...config.Option,
 ) (_ discovery.Client, err error) {
 	c := &client{
 		config:  config.New(opts...),
 		service: Ydb_Discovery_V1.NewDiscoveryServiceClient(cc),
 	}
 
+	crudExplorer.Lock()
+	defer crudExplorer.Unlock()
+
 	if c.config.Interval() <= 0 {
-		_ = crudExplorer.Insert(ctx, cc.Endpoint())
+		_ = crudExplorer.Insert(ctx, cc.Endpoint(), cluster.WithoutLock())
 		return c, nil
 	}
 
@@ -49,6 +52,7 @@ func New(
 		crudExplorer.Insert(
 			ctx,
 			e,
+			cluster.WithoutLock(),
 		)
 	}
 
@@ -70,25 +74,31 @@ func New(
 				// NOTE: curr endpoints must be sorted here.
 				cluster.SortEndpoints(next)
 
+				crudExplorer.Lock()
+				defer crudExplorer.Unlock()
+
 				cluster.DiffEndpoints(curr, next,
 					func(i, j int) {
 						// Endpoints are equal, but we still need to update meta
 						// data such that load factor and others.
 						crudExplorer.Update(
 							ctx,
 							next[j],
+							cluster.WithoutLock(),
 						)
 					},
 					func(i, j int) {
 						crudExplorer.Insert(
 							ctx,
 							next[j],
+							cluster.WithoutLock(),
 						)
 					},
 					func(i, j int) {
 						crudExplorer.Remove(
 							ctx,
 							curr[i],
+							cluster.WithoutLock(),
 						)
 					},
 				)

internal/meta/version.go

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

test/table_test.go

@@ -38,6 +38,13 @@ const (
 	folder = "table_test"
 )
 
+func init() {
+	log.SetFlags(0)
+	if os.Getenv("HIDE_APPLICATION_OUTPUT") == "1" {
+		log.SetOutput(ioutil.Discard)
+	}
+}
+
 type stats struct {
 	sync.Mutex
 
@@ -1018,11 +1025,6 @@ func days(date string) time.Time {
 	return t
 }
 
-func init() {
-	log.SetFlags(0)
-	log.SetOutput(ioutil.Discard)
-}
-
 type templateConfig struct {
 	TablePathPrefix string
 }