fix

Author: asmyasnikov

driver.go

@@ -4,11 +4,12 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"google.golang.org/grpc"
 	"os"
 	"sync"
 	"sync/atomic"
 
+	"google.golang.org/grpc"
+
 	"github.com/ydb-platform/ydb-go-sdk/v3/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/coordination"
 	"github.com/ydb-platform/ydb-go-sdk/v3/discovery"
@@ -102,7 +103,7 @@ type (
 
 		children    map[uint64]*Driver
 		childrenMtx xsync.Mutex
-		onClose     xsync.Map[uint64, func(c *Driver)]
+		onClose     []func(c *Driver)
 		closed      atomic.Bool
 
 		panicCallback func(e interface{})
@@ -163,11 +164,9 @@ func (d *Driver) Close(ctx context.Context) (finalErr error) {
 	d.ctxCancel()
 
 	defer func() {
-		d.onClose.Range(func(_ uint64, f func(c *Driver)) bool {
-			f(d)
-
-			return true
-		})
+		for _, onClose := range d.onClose {
+			onClose(d)
+		}
 	}()
 
 	closes := make([]func(context.Context) error, 0)

internal/xsql/connector.go

@@ -35,7 +35,6 @@ var (
 type (
 	Engine    uint8
 	Connector struct {
-		ID       uint64
 		parent   ydbDriver
 		balancer grpc.ClientConnInterface
 

options.go

@@ -832,7 +832,7 @@ func WithTraceDatabaseSQL(t trace.DatabaseSQL, opts ...trace.DatabaseSQLComposeO
 
 func withOnClose(onClose func(c *Driver)) Option {
 	return func(ctx context.Context, d *Driver) error {
-		d.onClose.Append(onClose)
+		d.onClose = append(d.onClose, onClose)
 
 		return nil
 	}

sql.go

@@ -5,14 +5,12 @@ import (
 	"database/sql"
 	"database/sql/driver"
 	"fmt"
-	"sync/atomic"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/bind"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/legacy"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/propose"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table/options"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
@@ -33,9 +31,7 @@ func withConnectorOptions(opts ...ConnectorOption) Option {
 	}
 }
 
-type sqlDriver struct {
-	connectors xsync.Map[*xsql.Connector, *Driver]
-}
+type sqlDriver struct{}
 
 var (
 	_ driver.Driver        = &sqlDriver{}
@@ -53,45 +49,18 @@ func (d *sqlDriver) OpenConnector(dataSourceName string) (driver.Connector, erro
 		return nil, xerrors.WithStackTrace(fmt.Errorf("failed to connect by data source name '%s': %w", dataSourceName, err))
 	}
 
-	c, err := Connector(db, db.databaseSQLOptions...)
+	c, err := Connector(db, append(db.databaseSQLOptions,
+		xsql.WithOnClose(func(connector *xsql.Connector) {
+			_ = db.Close(context.Background())
+		}),
+	)...)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(fmt.Errorf("failed to create connector: %w", err))
 	}
 
 	return c, nil
 }
 
-var globalConnectorCounter atomic.Uint64
-
-func (d *sqlDriver) attach(c *xsql.Connector, parent *Driver) {
-	c.ID = globalConnectorCounter.Add(1)
-
-	if parent != nil {
-		parent.onClose.Append(func(_ *Driver) {
-			_ = c.Close()
-		})
-	}
-
-	d.connectors.Set(c, parent)
-}
-
-func (d *sqlDriver) detach(c *xsql.Connector) {
-	var countConnectorsPerDriver int
-	parent, _ := d.connectors.Extract(c)
-	if parent != nil {
-		d.connectors.Range(func(key *xsql.Connector, value *Driver) bool {
-			if value == parent {
-				countConnectorsPerDriver++
-			}
-
-			return true
-		})
-		if countConnectorsPerDriver == 0 {
-			_ = parent.Close(context.Background())
-		}
-	}
-}
-
 type QueryMode int
 
 const (
@@ -255,7 +224,6 @@ func Connector(parent *Driver, opts ...ConnectorOption) (SQLConnector, error) {
 				parent.databaseSQLOptions,
 				opts...,
 			),
-			xsql.WithOnClose(d.detach),
 			xsql.WithTraceRetry(parent.config.TraceRetry()),
 			xsql.WithRetryBudget(parent.config.RetryBudget()),
 		)...,
@@ -264,8 +232,6 @@ func Connector(parent *Driver, opts ...ConnectorOption) (SQLConnector, error) {
 		return nil, xerrors.WithStackTrace(err)
 	}
 
-	d.attach(c, parent)
-
 	return c, nil
 }
 

tests/integration/helpers_test.go

@@ -473,11 +473,12 @@ func driverEngine(db *sql.DB) (engine xsql.Engine) {
 func simpleDetectGoroutineLeak(t *testing.T) {
 	// 1) testing.go => main.main()
 	// 2) current test
-	const expectedGoroutinesCount = 2
+	var expectedGoroutinesCount = 2
 	if num := runtime.NumGoroutine(); num > expectedGoroutinesCount {
 		bb := make([]byte, 2<<32)
 		if n := runtime.Stack(bb, true); n < len(bb) {
-			t.Error(fmt.Sprintf("unexpected goroutines:\n%s\n", string(bb[:n])))
+			bb = bb[:n]
 		}
+		t.Error(fmt.Sprintf("unexpected goroutines:\n%s\n", string(bb[runtime.Stack(bb, false)+1:])))
 	}
 }