* Implemented goroutine for closing idle connection in database/sql driver

Author: asmyasnikov

CHANGELOG.md

@@ -1,5 +1,6 @@
 * Fixed closing of `database/sql` connection (aka `YDB` session)
 * Made `session.Close()` as `nop` for idled session
+* Implemented goroutine for closing idle connection in `database/sql` driver
 
 ## v3.42.4
 * Added `ydb.WithDisableServerBalancer()` database/sql connector option

internal/xsql/conn.go

@@ -98,6 +98,7 @@ func newConn(c *Connector, s table.ClosableSession, opts ...connOption) *conn {
 	for _, o := range opts {
 		o(cc)
 	}
+	c.attach(cc)
 	return cc
 }
 
@@ -124,6 +125,10 @@ func (c *conn) PrepareContext(ctx context.Context, query string) (_ driver.Stmt,
 	}, nil
 }
 
+func (c *conn) sinceLastUsage() time.Duration {
+	return time.Since(time.Unix(atomic.LoadInt64(&c.lastUsage), 0))
+}
+
 func (c *conn) execContext(ctx context.Context, query string, args []driver.NamedValue) (_ driver.Result, err error) {
 	m := queryModeFromContext(ctx, c.defaultQueryMode)
 	onDone := trace.DatabaseSQLOnConnExec(
@@ -132,7 +137,7 @@ func (c *conn) execContext(ctx context.Context, query string, args []driver.Name
 		query,
 		m.String(),
 		xcontext.IsIdempotent(ctx),
-		time.Since(time.Unix(atomic.LoadInt64(&c.lastUsage), 0)),
+		c.sinceLastUsage(),
 	)
 	defer func() {
 		atomic.StoreInt64(&c.lastUsage, time.Now().Unix())
@@ -215,7 +220,7 @@ func (c *conn) queryContext(ctx context.Context, query string, args []driver.Nam
 		query,
 		m.String(),
 		xcontext.IsIdempotent(ctx),
-		time.Since(time.Unix(atomic.LoadInt64(&c.lastUsage), 0)),
+		c.sinceLastUsage(),
 	)
 	defer func() {
 		atomic.StoreInt64(&c.lastUsage, time.Now().Unix())
@@ -303,6 +308,7 @@ func (c *conn) Ping(ctx context.Context) (err error) {
 
 func (c *conn) Close() (err error) {
 	if atomic.CompareAndSwapUint32(&c.closed, 0, 1) {
+		c.connector.detach(c)
 		onDone := trace.DatabaseSQLOnConnClose(c.trace)
 		defer func() {
 			onDone(err)
@@ -313,7 +319,7 @@ func (c *conn) Close() (err error) {
 		}
 		return nil
 	}
-	return badconn.Map(xerrors.WithStackTrace(errClosedConn))
+	return badconn.Map(xerrors.WithStackTrace(errConnClosedEarly))
 }
 
 func (c *conn) Prepare(string) (driver.Stmt, error) {

internal/xsql/connector.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"database/sql/driver"
 	"io"
+	"sync"
+	"time"
 
 	metaHeaders "github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
@@ -58,10 +60,18 @@ func WithDisableServerBalancer() ConnectorOption {
 	}
 }
 
+func WithIdleThreshold(idleThreshold time.Duration) ConnectorOption {
+	return func(c *Connector) error {
+		c.idleThreshold = idleThreshold
+		return nil
+	}
+}
+
 func Open(d Driver, connection connection, opts ...ConnectorOption) (_ *Connector, err error) {
 	c := &Connector{
 		driver:           d,
 		connection:       connection,
+		conns:            make(map[*conn]struct{}),
 		defaultTxControl: table.DefaultTxControl(),
 		defaultQueryMode: DefaultQueryMode,
 	}
@@ -70,6 +80,9 @@ func Open(d Driver, connection connection, opts ...ConnectorOption) (_ *Connecto
 			return nil, err
 		}
 	}
+	if c.idleThreshold > 0 {
+		c.idleStopper = c.idleCloser()
+	}
 	d.Attach(c)
 	return c, nil
 }
@@ -97,11 +110,17 @@ type Connector struct {
 	driver     Driver
 	connection connection
 
+	conns    map[*conn]struct{}
+	connsMtx sync.RWMutex
+
+	idleStopper func()
+
 	defaultTxControl      *table.TransactionControl
 	defaultQueryMode      QueryMode
 	defaultDataQueryOpts  []options.ExecuteDataQueryOption
 	defaultScanQueryOpts  []options.ExecuteScanQueryOption
 	disableServerBalancer bool
+	idleThreshold         time.Duration
 
 	trace trace.DatabaseSQL
 }
@@ -111,15 +130,56 @@ var (
 	_ io.Closer        = &Connector{}
 )
 
+func (c *Connector) idleCloser() (idleStopper func()) {
+	var ctx context.Context
+	ctx, idleStopper = context.WithCancel(context.Background())
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-time.After(c.idleThreshold):
+				c.connsMtx.RLock()
+				conns := make([]*conn, 0, len(c.conns))
+				for cc := range c.conns {
+					conns = append(conns, cc)
+				}
+				c.connsMtx.RUnlock()
+				for _, cc := range conns {
+					if cc.sinceLastUsage() > c.idleThreshold {
+						cc.session.Close(context.Background())
+					}
+				}
+			}
+		}
+	}()
+	return idleStopper
+}
+
 func (c *Connector) Close() (err error) {
 	defer c.driver.Detach(c)
+	if c.idleStopper != nil {
+		c.idleStopper()
+	}
 	return nil
 }
 
 func (c *Connector) Connection() connection {
 	return c.connection
 }
 
+func (c *Connector) attach(cc *conn) {
+	c.connsMtx.Lock()
+	defer c.connsMtx.Unlock()
+	c.conns[cc] = struct{}{}
+}
+
+func (c *Connector) detach(cc *conn) {
+	c.connsMtx.Lock()
+	defer c.connsMtx.Unlock()
+	delete(c.conns, cc)
+}
+
 func (c *Connector) Connect(ctx context.Context) (_ driver.Conn, err error) {
 	onDone := trace.DatabaseSQLOnConnectorConnect(c.trace, &ctx)
 	defer func() {

internal/xsql/errors.go

@@ -8,8 +8,8 @@ import (
 )
 
 var (
-	ErrUnsupported  = driver.ErrSkip
-	errDeprecated   = driver.ErrSkip
-	errClosedConn   = xerrors.Retryable(errors.New("conn closed early"), xerrors.WithDeleteSession())
-	errNotReadyConn = xerrors.Retryable(errors.New("conn not ready"), xerrors.WithDeleteSession())
+	ErrUnsupported     = driver.ErrSkip
+	errDeprecated      = driver.ErrSkip
+	errConnClosedEarly = xerrors.Retryable(errors.New("conn closed early"), xerrors.WithDeleteSession())
+	errNotReadyConn    = xerrors.Retryable(errors.New("conn not ready"), xerrors.WithDeleteSession())
 )

options.go

@@ -360,6 +360,10 @@ func WithSessionPoolKeepAliveMinSize(keepAliveMinSize int) Option {
 func WithSessionPoolIdleThreshold(idleThreshold time.Duration) Option {
 	return func(ctx context.Context, c *connection) error {
 		c.tableOptions = append(c.tableOptions, tableConfig.WithIdleThreshold(idleThreshold))
+		c.databaseSQLOptions = append(
+			c.databaseSQLOptions,
+			xsql.WithIdleThreshold(idleThreshold),
+		)
 		return nil
 	}
 }