shared conn.Pool

Author: asmyasnikov

CHANGELOG.md

@@ -3,6 +3,7 @@
 * Removed `CustomOption` and taking client with custom options
 * Removed `proxy` package
 * Improved `db.With()` helper for create child connections
+* Setted shared `conn.Pool` for all child `ydb.Connection`  
 
 ## 3.12.1
 * Added `trace.Driver.OnConnPark` event

connection.go

@@ -83,6 +83,8 @@ type connection struct {
 	ratelimiter        ratelimiter.Client
 	ratelimiterOptions []ratelimiterConfig.Option
 
+	pool conn.Pool
+
 	mtx sync.Mutex
 	db  db.Connection
 
@@ -95,6 +97,12 @@ func (c *connection) Close(ctx context.Context) error {
 	c.mtx.Lock()
 	defer c.mtx.Unlock()
 
+	defer func() {
+		for _, f := range c.onClose {
+			f(c)
+		}
+	}()
+
 	var (
 		issues   []error
 		children = make([]Connection, 0, len(c.children))
@@ -136,12 +144,12 @@ func (c *connection) Close(ctx context.Context) error {
 		issues = append(issues, err)
 	}
 
-	if len(issues) > 0 {
-		return errors.NewWithIssues("close failed", issues...)
+	if err := c.pool.Release(ctx); err != nil {
+		issues = append(issues, err)
 	}
 
-	for _, f := range c.onClose {
-		f(c)
+	if len(issues) > 0 {
+		return errors.NewWithIssues("close failed", issues...)
 	}
 
 	return nil
@@ -257,9 +265,17 @@ func New(ctx context.Context, opts ...Option) (_ Connection, err error) {
 		)
 	}
 
+	if c.pool == nil {
+		c.pool = conn.NewPool(
+			ctx,
+			c.config,
+		)
+	}
+
 	c.db, err = db.New(
 		ctx,
 		c.config,
+		c.pool,
 		append(
 			// prepend config params from root config
 			[]discoveryConfig.Option{

internal/cluster/cluster.go

@@ -46,7 +46,7 @@ var (
 
 type cluster struct {
 	config   config.Config
-	pool     conn.PoolGetterCloser
+	pool     conn.Pool
 	dial     func(context.Context, string) (*grpc.ClientConn, error)
 	balancer balancer.Balancer
 	explorer repeater.Repeater
@@ -58,6 +58,47 @@ type cluster struct {
 	closed bool
 }
 
+func (c *cluster) Pessimize(ctx context.Context, cc conn.Conn, cause error) {
+	c.pool.Pessimize(ctx, cc, cause)
+
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	if c.closed {
+		return
+	}
+
+	entry, has := c.index[cc.Endpoint().Address()]
+	if !has {
+		return
+	}
+
+	if entry.Handle == nil {
+		return
+	}
+
+	if !c.balancer.Contains(entry.Handle) {
+		return
+	}
+
+	if c.explorer == nil {
+		return
+	}
+
+	// count ratio (banned/all)
+	online := 0
+	for _, entry = range c.index {
+		if entry.Conn != nil && entry.Conn.GetState() == conn.Online {
+			online++
+		}
+	}
+
+	// more than half connections banned - re-discover now
+	if online*2 < len(c.index) {
+		c.explorer.Force()
+	}
+}
+
 func (c *cluster) Lock() {
 	c.mu.Lock()
 }
@@ -136,75 +177,34 @@ type Cluster interface {
 	Explorer
 	Locker
 	conn.PoolGetter
+	conn.Pessimizer
 }
 
 func New(
 	ctx context.Context,
 	config config.Config,
+	pool conn.Pool,
 	balancer balancer.Balancer,
 ) Cluster {
 	onDone := trace.DriverOnClusterInit(config.Trace(), &ctx)
 	defer func() {
-		onDone()
+		onDone(pool.Take(ctx))
 	}()
 
-	c := &cluster{
+	return &cluster{
 		config:    config,
 		index:     make(map[string]entry.Entry),
 		endpoints: make(map[uint32]conn.Conn),
+		pool:      pool,
 		balancer:  balancer,
 	}
-
-	c.pool = conn.NewPool(
-		ctx,
-		config,
-		func(e endpoint.Endpoint) {
-			c.mu.RLock()
-			defer c.mu.RUnlock()
-
-			if c.closed {
-				return
-			}
-
-			entry, has := c.index[e.Address()]
-			if !has {
-				return
-			}
-
-			if entry.Handle == nil {
-				return
-			}
-
-			if !c.balancer.Contains(entry.Handle) {
-				return
-			}
-
-			if c.explorer == nil {
-				return
-			}
-
-			// count ratio (banned/all)
-			online := 0
-			for _, entry = range c.index {
-				if entry.Conn != nil && entry.Conn.GetState() == conn.Online {
-					online++
-				}
-			}
-
-			// more than half connections banned - re-discover now
-			if online*2 < len(c.index) {
-				c.explorer.Force()
-			}
-		},
-	)
-
-	return c
 }
 
 func (c *cluster) Close(ctx context.Context) (err error) {
 	c.mu.Lock()
+	defer c.mu.Unlock()
+
 	if c.closed {
-		c.mu.Unlock()
 		return
 	}
 
@@ -218,19 +218,10 @@ func (c *cluster) Close(ctx context.Context) (err error) {
 	}
 	c.closed = true
 
-	index := c.index
 	c.index = nil
 	c.endpoints = nil
 
-	c.mu.Unlock()
-
-	for _, entry := range index {
-		if entry.Conn != nil {
-			_ = entry.Conn.Close(ctx)
-		}
-	}
-
-	return c.pool.Close(ctx)
+	return c.pool.Release(ctx)
 }
 
 // Get returns next available connection.

internal/cluster/cluster_test.go

@@ -39,7 +39,7 @@ func TestClusterFastRedial(t *testing.T) {
 		balancer:  b,
 		index:     make(map[string]entry.Entry),
 		endpoints: make(map[uint32]conn.Conn),
-		pool:      conn.NewPool(ctx, config.New(), nil),
+		pool:      conn.NewPool(ctx, config.New()),
 	}
 
 	pingConnects := func(size int) chan struct{} {
@@ -96,7 +96,7 @@ func TestClusterMergeEndpoints(t *testing.T) {
 		}(),
 		index:     make(map[string]entry.Entry),
 		endpoints: make(map[uint32]conn.Conn),
-		pool:      conn.NewPool(ctx, config.New(), nil),
+		pool:      conn.NewPool(ctx, config.New()),
 	}
 
 	assert := func(t *testing.T, exp []endpoint.Endpoint) {

internal/conn/conn.go

@@ -24,12 +24,14 @@ type Conn interface {
 
 	Endpoint() endpoint.Endpoint
 
+	TTL() <-chan time.Time
+
 	IsState(states ...State) bool
 	GetState() State
 	SetState(State) State
+
 	Close(ctx context.Context) error
 	Park(ctx context.Context) error
-	TTL() <-chan time.Time
 }
 
 func (c *conn) Address() string {
@@ -38,16 +40,15 @@ func (c *conn) Address() string {
 
 type conn struct {
 	sync.RWMutex
-	config      Config // ro access
-	cc          *grpc.ClientConn
-	done        chan struct{}
-	endpoint    endpoint.Endpoint // ro access
-	closed      bool
-	state       State
-	usages      int32
-	ttl         timeutil.Timer
-	onClose     []func(Conn)
-	onPessimize []func(e endpoint.Endpoint)
+	config   Config // ro access
+	cc       *grpc.ClientConn
+	done     chan struct{}
+	endpoint endpoint.Endpoint // ro access
+	closed   bool
+	state    State
+	usages   int32
+	ttl      timeutil.Timer
+	onClose  []func(Conn)
 }
 
 func (c *conn) IsState(states ...State) bool {
@@ -262,39 +263,13 @@ func (c *conn) Close(ctx context.Context) (err error) {
 	return err
 }
 
-func (c *conn) pessimize(ctx context.Context, err error) {
-	if c.isClosed() {
-		return
-	}
-
-	defer func() {
-		for _, f := range c.onPessimize {
-			f(c.endpoint)
-		}
-	}()
-
-	trace.DriverOnPessimizeNode(
-		trace.ContextDriver(ctx).Compose(c.config.Trace()),
-		&ctx,
-		c.endpoint.Copy(),
-		c.GetState(),
-		err,
-	)(c.SetState(Banned))
-}
-
 func (c *conn) invoke(
 	ctx context.Context,
 	method string,
 	req interface{},
 	res interface{},
 	opts ...grpc.CallOption,
 ) (err error) {
-	defer func() {
-		if err != nil && errors.MustPessimizeEndpoint(err) {
-			c.pessimize(ctx, err)
-		}
-	}()
-
 	var cc *grpc.ClientConn
 	cc, err = c.take(ctx)
 	if err != nil {
@@ -368,12 +343,6 @@ func (c *conn) newStream(
 	method string,
 	opts ...grpc.CallOption,
 ) (_ grpc.ClientStream, err error) {
-	defer func() {
-		if err != nil && errors.MustPessimizeEndpoint(err) {
-			c.pessimize(ctx, err)
-		}
-	}()
-
 	var cc *grpc.ClientConn
 	cc, err = c.take(ctx)
 	if err != nil {
@@ -449,14 +418,6 @@ func (c *conn) NewStream(
 
 type option func(c *conn)
 
-func withOnPessimize(onPessimize func(e endpoint.Endpoint)) option {
-	return func(c *conn) {
-		if onPessimize != nil {
-			c.onPessimize = append(c.onPessimize, onPessimize)
-		}
-	}
-}
-
 func withOnClose(onClose func(Conn)) option {
 	return func(c *conn) {
 		if onClose != nil {
@@ -467,12 +428,11 @@ func withOnClose(onClose func(Conn)) option {
 
 func New(e endpoint.Endpoint, config Config, opts ...option) Conn {
 	c := &conn{
-		state:       Created,
-		endpoint:    e,
-		config:      config,
-		done:        make(chan struct{}),
-		onClose:     make([]func(Conn), 0),
-		onPessimize: make([]func(e endpoint.Endpoint), 0),
+		state:    Created,
+		endpoint: e,
+		config:   config,
+		done:     make(chan struct{}),
+		onClose:  make([]func(Conn), 0),
 	}
 	for _, o := range opts {
 		o(c)

internal/conn/grpc_client_stream.go

@@ -66,9 +66,6 @@ func (s *grpcClientStream) RecvMsg(m interface{}) (err error) {
 	err = s.s.RecvMsg(m)
 
 	if err != nil {
-		if errors.MustPessimizeEndpoint(err) {
-			s.c.pessimize(s.s.Context(), err)
-		}
 		if s.wrapping {
 			return errors.Errorf(0, "receive message failed: %w", errors.MapGRPCError(err))
 		}