fix: data-race with concurrent scan and close (#304)

* fix: data-race with concurrent scan and close Fixes #299 * test: add test for race condition in close --------- Co-authored-by: Knut Olav Løite <[email protected]>
googleapis · Oct 16, 2024 · cdbea5c · cdbea5c
1 parent 081fa51
commit cdbea5c
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 5 deletions.
diff --git a/driver.go b/driver.go
@@ -171,7 +171,7 @@ type connector struct {
 	// propagated to the caller. This option is enabled by default.
 	retryAbortsInternally bool
 
-	initClient     sync.Mutex
+	clientMu       sync.Mutex
 	client         *spanner.Client
 	clientErr      error
 	adminClient    *adminapi.DatabaseAdminClient
@@ -317,8 +317,8 @@ func openDriverConn(ctx context.Context, c *connector) (driver.Conn, error) {
 
 // increaseConnCount initializes the client and increases the number of connections that are active.
 func (c *connector) increaseConnCount(ctx context.Context, databaseName string, opts []option.ClientOption) error {
-	c.initClient.Lock()
-	defer c.initClient.Unlock()
+	c.clientMu.Lock()
+	defer c.clientMu.Unlock()
 
 	if c.clientErr != nil {
 		return c.clientErr
@@ -349,8 +349,8 @@ func (c *connector) increaseConnCount(ctx context.Context, databaseName string,
 // decreaseConnCount decreases the number of connections that are active and closes the underlying clients if it was the
 // last connection.
 func (c *connector) decreaseConnCount() error {
-	c.initClient.Lock()
-	defer c.initClient.Unlock()
+	c.clientMu.Lock()
+	defer c.clientMu.Unlock()
 
 	c.connCount--
 	if c.connCount > 0 {
@@ -373,6 +373,8 @@ func (c *connector) Close() error {
 	delete(c.driver.connectors, c.dsn)
 	c.driver.mu.Unlock()
 
+	c.clientMu.Lock()
+	defer c.clientMu.Unlock()
 	return c.closeClients()
 }
 

diff --git a/driver_with_mockserver_test.go b/driver_with_mockserver_test.go
@@ -119,6 +119,39 @@ func TestSimpleQuery(t *testing.T) {
 	}
 }
 
+func TestConcurrentScanAndClose(t *testing.T) {
+	t.Parallel()
+
+	db, _, teardown := setupTestDBConnection(t)
+	defer teardown()
+	rows, err := db.QueryContext(context.Background(), testutil.SelectFooFromBar)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Only fetch the first row of the query to make sure that the rows are not auto-closed
+	// when the end of the stream is reached.
+	rows.Next()
+	var got int64
+	err = rows.Scan(&got)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Close both the database and the rows (connection) in parallel.
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		_ = db.Close()
+	}()
+	go func() {
+		defer wg.Done()
+		_ = rows.Close()
+	}()
+	wg.Wait()
+}
+
 func TestSingleQueryWithTimestampBound(t *testing.T) {
 	t.Parallel()