Merge pull request #1587 from ydb-platform/issue-1585

fixed bug with multiple closing driver

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed panic on multiple closing driver
+
 ## v3.95.1
 * Added alias from `ydb.WithFakeTx(ydb.ScriptingQueryMode)` to `ydb.WithFakeTx(ydb.QueryExecuteQueryMode)` for compatibility with legacy code   
 

driver.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 	"sync"
+	"sync/atomic"
 
 	"google.golang.org/grpc"
 
@@ -103,6 +104,7 @@ type (
 		children    map[uint64]*Driver
 		childrenMtx xsync.Mutex
 		onClose     []func(c *Driver)
+		closed      atomic.Bool
 
 		panicCallback func(e interface{})
 	}
@@ -149,6 +151,11 @@ func (d *Driver) Close(ctx context.Context) (finalErr error) {
 	defer func() {
 		onDone(finalErr)
 	}()
+
+	if !d.closed.CompareAndSwap(false, true) {
+		return nil
+	}
+
 	d.ctxCancel()
 
 	d.mtx.Lock()

internal/query/client.go

@@ -188,6 +188,10 @@ func (c *Client) ExecuteScript(
 }
 
 func (c *Client) Close(ctx context.Context) error {
+	if c == nil {
+		return xerrors.WithStackTrace(errNilClient)
+	}
+
 	close(c.done)
 
 	if err := c.pool.Close(ctx); err != nil {

internal/query/errors.go

@@ -7,6 +7,7 @@ import (
 )
 
 var (
+	errNilClient               = xerrors.Wrap(errors.New("table client is not initialized"))
 	ErrTransactionRollingBack  = xerrors.Wrap(errors.New("ydb: the transaction is rolling back"))
 	errWrongNextResultSetIndex = errors.New("wrong result set index")
 	errWrongResultSetIndex     = errors.New("critical violation of the logic - wrong result set index")

tests/integration/driver_test.go

@@ -651,3 +651,22 @@ func TestClusterDiscoveryRetry(t *testing.T) {
 	}
 	require.Greater(t, counter, 1)
 }
+
+func TestMultipleClosingDriverIssue1585(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	db, err := ydb.Open(ctx, os.Getenv("YDB_CONNECTION_STRING"))
+	require.NoError(t, err)
+
+	require.NotPanics(t, func() {
+		err = db.Close(ctx)
+		require.NoError(t, err)
+
+		err = db.Close(ctx)
+		require.NoError(t, err)
+
+		err = db.Close(ctx)
+		require.NoError(t, err)
+	})
+}