Merge pull request #252 from ydb-platform/err-cluster-closed

* Fixed bug with unexpected failing of call `Invoke` and `NewStream` …

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,7 @@
+* Fixed bug with unexpected failing of call `Invoke` and `NewStream` on closed cluster
+* Fixed bug with releasing `internal/conn/conn.Pool` in cluster
+* Replaced interface `internal/conn/conn.Pool` to struct `internal/conn/conn.Pool`
+
 ## v3.25.0
 * Added `ydb.GRPCConn(ydb.Connection)` helper for connect to driver-unsupported YDB services
 * Marked as deprecated `session.Prepare` callback

connection.go

@@ -102,7 +102,7 @@ type connection struct {
 	ratelimiter        *internalRatelimiter.Client
 	ratelimiterOptions []ratelimiterConfig.Option
 
-	pool conn.Pool
+	pool *conn.Pool
 
 	mtx    sync.Mutex
 	router router.Connection

internal/cluster/cluster.go

@@ -19,38 +19,26 @@ const (
 	MaxGetConnTimeout = 10 * time.Second
 )
 
+// nolint: gofumpt
+// nolint: nolintlint
 var (
-	// ErrClusterClosed returned when requested on a closed cluster.
-	ErrClusterClosed = xerrors.Wrap(fmt.Errorf("cluster closed"))
-
 	// ErrClusterEmpty returned when no connections left in cluster.
 	ErrClusterEmpty = xerrors.Wrap(fmt.Errorf("cluster empty"))
 )
 
 type Cluster struct {
 	config config.Config
-	pool   conn.Pool
+	pool   *conn.Pool
 	conns  []conn.Conn
 
 	balancer balancer.Balancer
 
-	done chan struct{}
-
 	onBadStateCallback balancer.OnBadStateCallback
 }
 
-func (c *Cluster) isClosed() bool {
-	select {
-	case <-c.done:
-		return true
-	default:
-		return false
-	}
-}
-
 // Ban connection in underling pool
 func (c *Cluster) Ban(ctx context.Context, cc conn.Conn, cause error) {
-	c.pool.Pessimize(ctx, cc, cause)
+	c.pool.Ban(ctx, cc, cause)
 
 	online := 0
 	for _, cc := range c.conns {
@@ -64,15 +52,15 @@ func (c *Cluster) Ban(ctx context.Context, cc conn.Conn, cause error) {
 	}
 }
 
-// Unban connection in underling pool
-func (c *Cluster) Unban(ctx context.Context, cc conn.Conn) {
-	c.pool.Unpessimize(ctx, cc)
+// Allow connection in underling pool
+func (c *Cluster) Allow(ctx context.Context, cc conn.Conn) {
+	c.pool.Allow(ctx, cc)
 }
 
 func New(
 	ctx context.Context,
 	config config.Config,
-	pool conn.Pool,
+	pool *conn.Pool,
 	endpoints []endpoint.Endpoint,
 	onBadStateCallback balancer.OnBadStateCallback,
 ) *Cluster {
@@ -99,7 +87,6 @@ func New(
 	)
 
 	return &Cluster{
-		done:               make(chan struct{}),
 		config:             config,
 		pool:               pool,
 		conns:              conns,
@@ -118,17 +105,12 @@ func (c *Cluster) Close(ctx context.Context) (err error) {
 		onDone(err)
 	}()
 
-	close(c.done)
-
-	return nil
+	return c.pool.Release(ctx)
 }
 
 func (c *Cluster) get(ctx context.Context) (cc conn.Conn, _ error) {
 	for {
 		select {
-		case <-c.done:
-			return nil, xerrors.WithStackTrace(ErrClusterClosed)
-
 		case <-ctx.Done():
 			return nil, xerrors.WithStackTrace(ctx.Err())
 
@@ -156,10 +138,6 @@ func (c *Cluster) get(ctx context.Context) (cc conn.Conn, _ 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) {
-	if c.isClosed() {
-		return nil, xerrors.WithStackTrace(ErrClusterClosed)
-	}
-
 	var cancel context.CancelFunc
 	// without client context deadline lock limited on MaxGetConnTimeout
 	// cluster endpoints cannot be updated at this time

internal/conn/pool.go

@@ -14,36 +14,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
-type Pool interface {
-	Getter
-	Taker
-	Releaser
-	Pessimizer
-}
-
-type Getter interface {
-	Get(endpoint endpoint.Endpoint) Conn
-}
-
-type Taker interface {
-	Take(ctx context.Context) error
-}
-
-type Releaser interface {
-	Release(ctx context.Context) error
-}
-
-type Pessimizer interface {
-	Pessimize(ctx context.Context, cc Conn, cause error)
-	Unpessimize(ctx context.Context, cc Conn)
-}
-
-type PoolConfig interface {
-	ConnectionTTL() time.Duration
-	GrpcDialOptions() []grpc.DialOption
-}
-
-type pool struct {
+type Pool struct {
 	usages int64
 	config Config
 	mtx    sync.RWMutex
@@ -52,7 +23,7 @@ type pool struct {
 	done   chan struct{}
 }
 
-func (p *pool) Get(endpoint endpoint.Endpoint) Conn {
+func (p *Pool) Get(endpoint endpoint.Endpoint) Conn {
 	p.mtx.Lock()
 	defer p.mtx.Unlock()
 
@@ -73,13 +44,13 @@ func (p *pool) Get(endpoint endpoint.Endpoint) Conn {
 	return cc
 }
 
-func (p *pool) remove(c *conn) {
+func (p *Pool) remove(c *conn) {
 	p.mtx.Lock()
 	defer p.mtx.Unlock()
 	delete(p.conns, c.Endpoint().Address())
 }
 
-func (p *pool) Pessimize(ctx context.Context, cc Conn, cause error) {
+func (p *Pool) Ban(ctx context.Context, cc Conn, cause error) {
 	e := cc.Endpoint().Copy()
 
 	p.mtx.RLock()
@@ -90,7 +61,7 @@ func (p *pool) Pessimize(ctx context.Context, cc Conn, cause error) {
 		return
 	}
 
-	trace.DriverOnPessimizeNode(
+	trace.DriverOnConnBan(
 		p.config.Trace(),
 		&ctx,
 		e,
@@ -99,7 +70,7 @@ func (p *pool) Pessimize(ctx context.Context, cc Conn, cause error) {
 	)(cc.SetState(Banned))
 }
 
-func (p *pool) Unpessimize(ctx context.Context, cc Conn) {
+func (p *Pool) Allow(ctx context.Context, cc Conn) {
 	e := cc.Endpoint().Copy()
 
 	p.mtx.RLock()
@@ -110,19 +81,19 @@ func (p *pool) Unpessimize(ctx context.Context, cc Conn) {
 		return
 	}
 
-	trace.DriverOnUnpessimizeNode(p.config.Trace(),
+	trace.DriverOnConnAllow(p.config.Trace(),
 		&ctx,
 		e,
 		cc.GetState(),
 	)(cc.SetState(Online))
 }
 
-func (p *pool) Take(context.Context) error {
+func (p *Pool) Take(context.Context) error {
 	atomic.AddInt64(&p.usages, 1)
 	return nil
 }
 
-func (p *pool) Release(ctx context.Context) error {
+func (p *Pool) Release(ctx context.Context) error {
 	if atomic.AddInt64(&p.usages, -1) > 0 {
 		return nil
 	}
@@ -150,7 +121,7 @@ func (p *pool) Release(ctx context.Context) error {
 	return nil
 }
 
-func (p *pool) connParker(ctx context.Context, ttl, interval time.Duration) {
+func (p *Pool) connParker(ctx context.Context, ttl, interval time.Duration) {
 	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
 	for {
@@ -172,7 +143,7 @@ func (p *pool) connParker(ctx context.Context, ttl, interval time.Duration) {
 	}
 }
 
-func (p *pool) collectConns() []*conn {
+func (p *Pool) collectConns() []*conn {
 	p.mtx.RLock()
 	defer p.mtx.RUnlock()
 	conns := make([]*conn, 0, len(p.conns))
@@ -185,8 +156,8 @@ func (p *pool) collectConns() []*conn {
 func NewPool(
 	ctx context.Context,
 	config Config,
-) Pool {
-	p := &pool{
+) *Pool {
+	p := &Pool{
 		usages: 1,
 		config: config,
 		opts:   config.GrpcDialOptions(),

internal/router/router.go

@@ -22,7 +22,7 @@ import (
 
 type router struct {
 	config config.Config
-	pool   conn.Pool
+	pool   *conn.Pool
 
 	clusterMtx sync.RWMutex
 	clusterPtr *cluster.Cluster
@@ -99,7 +99,7 @@ func (r *router) Close(ctx context.Context) (err error) {
 func New(
 	ctx context.Context,
 	c config.Config,
-	pool conn.Pool,
+	pool *conn.Pool,
 	opts ...discoveryConfig.Option,
 ) (_ Connection, err error) {
 	onDone := trace.DriverOnInit(
@@ -182,7 +182,17 @@ func (r *router) Invoke(
 		return xerrors.WithStackTrace(err)
 	}
 
-	defer r.handleConnRequestError(ctx, &err, cc)
+	defer func() {
+		if err == nil {
+			if cc.GetState() == conn.Banned {
+				r.cluster().Allow(ctx, cc)
+			}
+		} else {
+			if xerrors.MustPessimizeEndpoint(err, r.config.ExcludeGRPCCodesForPessimization()...) {
+				r.cluster().Ban(ctx, cc, err)
+			}
+		}
+	}()
 
 	ctx, err = r.config.Meta().Meta(ctx)
 	if err != nil {
@@ -202,13 +212,24 @@ func (r *router) NewStream(
 	desc *grpc.StreamDesc,
 	method string,
 	opts ...grpc.CallOption,
-) (grpc.ClientStream, error) {
-	cc, err := r.cluster().Get(ctx)
+) (_ grpc.ClientStream, err error) {
+	var cc conn.Conn
+	cc, err = r.cluster().Get(ctx)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}
 
-	defer r.handleConnRequestError(ctx, &err, cc)
+	defer func() {
+		if err == nil {
+			if cc.GetState() == conn.Banned {
+				r.cluster().Allow(ctx, cc)
+			}
+		} else {
+			if xerrors.MustPessimizeEndpoint(err, r.config.ExcludeGRPCCodesForPessimization()...) {
+				r.cluster().Ban(ctx, cc, err)
+			}
+		}
+	}()
 
 	ctx, err = r.config.Meta().Meta(ctx)
 	if err != nil {
@@ -223,13 +244,3 @@ func (r *router) NewStream(
 
 	return client, nil
 }
-
-func (r *router) handleConnRequestError(ctx context.Context, perr *error, cc conn.Conn) {
-	err := *perr
-	if err == nil && cc.GetState() == conn.Banned {
-		r.cluster().Unban(ctx, cc)
-	}
-	if err != nil && xerrors.MustPessimizeEndpoint(err, r.config.ExcludeGRPCCodesForPessimization()...) {
-		r.cluster().Ban(ctx, cc, err)
-	}
-}

internal/table/client.go

@@ -645,7 +645,6 @@ func (c *Client) keeper(ctx context.Context) {
 					case
 						xerrors.Is(
 							err,
-							cluster.ErrClusterClosed,
 							cluster.ErrClusterEmpty,
 						),
 						xerrors.IsOperationError(err, Ydb.StatusIds_BAD_SESSION),