@@ -45,6 +45,7 @@ import (
45
45
"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
46
46
"github.com/cockroachdb/cockroach/pkg/testutils/skip"
47
47
"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
48
+ "github.com/cockroachdb/cockroach/pkg/util/ctxgroup"
48
49
"github.com/cockroachdb/datadriven"
49
50
"github.com/cockroachdb/errors"
50
51
"github.com/lib/pq"
@@ -91,10 +92,9 @@ type sqlDBKey struct {
91
92
92
93
type datadrivenTestState struct {
93
94
// cluster maps the user defined cluster name to its cluster
94
- clusters map [string ]serverutils.TestClusterInterface
95
+ clusters map [string ]serverutils.TestClusterInterface
96
+ clusterCleanup []func ()
95
97
96
- // firstNode maps the cluster name to the first node in the cluster
97
- firstNode map [string ]serverutils.TestServerInterface
98
98
dataDirs map [string ]string
99
99
sqlDBs map [sqlDBKey ]* gosql.DB
100
100
jobTags map [string ]jobspb.JobID
@@ -107,7 +107,6 @@ type datadrivenTestState struct {
107
107
func newDatadrivenTestState () datadrivenTestState {
108
108
return datadrivenTestState {
109
109
clusters : make (map [string ]serverutils.TestClusterInterface ),
110
- firstNode : make (map [string ]serverutils.TestServerInterface ),
111
110
dataDirs : make (map [string ]string ),
112
111
sqlDBs : make (map [sqlDBKey ]* gosql.DB ),
113
112
jobTags : make (map [string ]jobspb.JobID ),
@@ -120,16 +119,27 @@ func (d *datadrivenTestState) cleanup(ctx context.Context, t *testing.T) {
120
119
// While the testCluster cleanupFns would close the dbConn and clusters, close
121
120
// them manually to ensure all queries finish on tests that share these
122
121
// resources.
122
+ for _ , f := range d .cleanupFns {
123
+ f ()
124
+ }
125
+
123
126
for _ , db := range d .sqlDBs {
124
127
backuptestutils .CheckForInvalidDescriptors (t , db )
125
128
db .Close ()
126
129
}
127
- for _ , s := range d .firstNode {
128
- s .Stopper ().Stop (ctx )
129
- }
130
- for _ , f := range d .cleanupFns {
131
- f ()
130
+
131
+ // Clean up the clusters in parallel so that if one gets stuck we don't see
132
+ // goroutines for running clusters mixed in with goroutines for the stuck
133
+ // cluster.
134
+ group := ctxgroup .WithContext (ctx )
135
+ for _ , cleanup := range d .clusterCleanup {
136
+ group .Go (func () error {
137
+ cleanup ()
138
+ return nil
139
+ })
132
140
}
141
+ require .NoError (t , group .Wait ())
142
+
133
143
d .noticeBuffer = nil
134
144
}
135
145
@@ -223,9 +233,8 @@ func (d *datadrivenTestState) addCluster(t *testing.T, cfg clusterCfg) error {
223
233
var cleanup func ()
224
234
tc , _ , cfg .iodir , cleanup = backuptestutils .StartBackupRestoreTestCluster (t , clusterSize , opts ... )
225
235
d .clusters [cfg .name ] = tc
226
- d .firstNode [cfg .name ] = tc .Server (0 )
227
236
d .dataDirs [cfg .name ] = cfg .iodir
228
- d .cleanupFns = append (d .cleanupFns , cleanup )
237
+ d .clusterCleanup = append (d .clusterCleanup , cleanup )
229
238
230
239
return nil
231
240
}
@@ -262,15 +271,15 @@ func (d *datadrivenTestState) getSQLDBForVC(
262
271
serverutils .User (user ),
263
272
}
264
273
265
- s := d .firstNode [name ].ApplicationLayer ()
274
+ s := d .clusters [name ].ApplicationLayer (0 )
266
275
switch vc {
267
276
case "default" :
268
277
// Nothing to do.
269
278
case "system" :
270
279
// We use the system layer since in the case of
271
280
// external SQL server's the application layer can't
272
281
// route to the system tenant.
273
- s = d .firstNode [name ].SystemLayer ()
282
+ s = d .clusters [name ].SystemLayer (0 )
274
283
default :
275
284
opts = append (opts , serverutils .DBName ("cluster:" + vc ))
276
285
}
@@ -956,7 +965,7 @@ func runTestDataDriven(t *testing.T, testFilePathFromWorkspace string) {
956
965
return ""
957
966
958
967
case "create-dummy-system-table" :
959
- al := ds .firstNode [lastCreatedCluster ].ApplicationLayer ()
968
+ al := ds .clusters [lastCreatedCluster ].ApplicationLayer (0 )
960
969
db := al .DB ()
961
970
execCfg := al .ExecutorConfig ().(sql.ExecutorConfig )
962
971
codec := execCfg .Codec
@@ -1037,7 +1046,7 @@ func handleKVRequest(
1037
1046
},
1038
1047
UseRangeTombstone : true ,
1039
1048
}
1040
- if _ , err := kv .SendWrapped (ctx , ds .firstNode [cluster ].SystemLayer ().DistSenderI ().(* kvcoord.DistSender ), & dr ); err != nil {
1049
+ if _ , err := kv .SendWrapped (ctx , ds .clusters [cluster ].SystemLayer (0 ).DistSenderI ().(* kvcoord.DistSender ), & dr ); err != nil {
1041
1050
t .Fatal (err )
1042
1051
}
1043
1052
} else {
0 commit comments