From e6e00c49a5a9d666e3c1aafd0ce286029111ec60 Mon Sep 17 00:00:00 2001
From: Marius Posta <marius@cockroachlabs.com>
Date: Mon, 17 Jul 2023 14:48:51 -0400
Subject: [PATCH] schemachanger: fix bug in concurrency test framework

The newly-rewritten TestConcurrentDeclarativeSchemaChanges test had
a concurrency bug in it, which caused it to fail in cases where the
two schema changes actually take place serially. It's possible for the
later one to start after the earlier one has already finished.

Fixes: #106732.

Release note: None
---
 pkg/sql/schemachanger/schemachanger_test.go | 36 +++++++++------------
 1 file changed, 16 insertions(+), 20 deletions(-)

diff --git a/pkg/sql/schemachanger/schemachanger_test.go b/pkg/sql/schemachanger/schemachanger_test.go
index e1f99c5dc787..8f3d75af572f 100644
--- a/pkg/sql/schemachanger/schemachanger_test.go
+++ b/pkg/sql/schemachanger/schemachanger_test.go
@@ -140,32 +140,28 @@ func TestConcurrentDeclarativeSchemaChanges(t *testing.T) {
 	}
 
 	backfillNotif, continueNotif := initBackfillNotification()
-	var wg sync.WaitGroup
-	wg.Add(1)
-	go func() {
-		if _, err := sqlDB.Exec(`CREATE INDEX i ON t.test (v)`); err != nil {
-			t.Error(err)
-		}
-		wg.Done()
-	}()
+	g := ctxgroup.WithContext(ctx)
+	g.Go(func() error {
+		_, err := sqlDB.Exec(`CREATE INDEX i ON t.test (v)`)
+		return err
+	})
 
-	// Wait until the create index schema change job has started
-	// before kicking off the alter primary key.
+	// Wait until the CREATE INDEX schema change job has started
+	// before kicking off the ALTER PRIMARY KEY.
 	<-backfillNotif
 	require.Zero(t, alterPrimaryKeyBlockedCounter.Load())
-	wg.Add(1)
-	go func() {
-		if _, err := sqlDB.Exec(`ALTER TABLE t.test ALTER PRIMARY KEY USING COLUMNS (k)`); err != nil {
-			t.Error(err)
-		}
-		wg.Done()
-	}()
+	g.Go(func() error {
+		_, err := sqlDB.Exec(`ALTER TABLE t.test ALTER PRIMARY KEY USING COLUMNS (k)`)
+		return err
+	})
 
-	// Unblock the create index job.
+	// Unblock the CREATE INDEX, and, by extension, ALTER PRIMARY KEY.
 	continueNotif <- struct{}{}
-	wg.Wait()
 
-	// The ALTER PRIMARY KEY schema change must have been blocked.
+	// Wait for both schema changes to complete.
+	require.NoError(t, g.Wait())
+
+	// The ALTER PRIMARY KEY schema change must have blocked.
 	require.NotZero(t, alterPrimaryKeyBlockedCounter.Load())
 
 	// There should be 5 k/v pairs per row: