@@ -8,12 +8,12 @@ package main
8
8
import (
9
9
"context"
10
10
"fmt"
11
+ "io"
11
12
"math"
12
13
"math/rand"
13
14
"os"
14
15
"os/user"
15
16
"regexp"
16
- "sort"
17
17
"strings"
18
18
"time"
19
19
@@ -25,6 +25,7 @@ import (
25
25
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/testselector"
26
26
"github.com/cockroachdb/cockroach/pkg/roachprod"
27
27
"github.com/cockroachdb/cockroach/pkg/roachprod/config"
28
+ "github.com/cockroachdb/cockroach/pkg/util/randutil"
28
29
"github.com/cockroachdb/errors"
29
30
_ "github.com/lib/pq" // register postgres driver
30
31
"github.com/spf13/cobra"
@@ -315,7 +316,12 @@ func testsToRun(
315
316
}
316
317
}
317
318
318
- return selectSpecs (notSkipped , selectProbability , true , print ), nil
319
+ var stdout io.Writer
320
+ if print {
321
+ stdout = os .Stdout
322
+ }
323
+ rng , _ := randutil .NewPseudoRand ()
324
+ return selectSpecs (notSkipped , rng , selectProbability , true , stdout ), nil
319
325
}
320
326
321
327
// updateSpecForSelectiveTests is responsible for updating the test spec skip and skip details
@@ -436,24 +442,28 @@ func opsToRun(r testRegistryImpl, filter string) ([]registry.OperationSpec, erro
436
442
// testRegistryImpl.AllTests().
437
443
// TODO(smg260): Perhaps expose `atLeastOnePerPrefix` via CLI
438
444
func selectSpecs (
439
- specs []registry.TestSpec , samplePct float64 , atLeastOnePerPrefix bool , print bool ,
445
+ specs []registry.TestSpec ,
446
+ rng * rand.Rand ,
447
+ samplePct float64 ,
448
+ atLeastOnePerPrefix bool ,
449
+ stdout io.Writer ,
440
450
) []registry.TestSpec {
441
451
if samplePct == 1 || len (specs ) == 0 {
442
452
return specs
443
453
}
444
454
445
455
var sampled []registry.TestSpec
446
- var selectedIdxs [] int
456
+ selectedIndexes := make ( map [ int ] struct {})
447
457
448
458
prefix := strings .Split (specs [0 ].Name , "/" )[0 ]
449
459
prefixSelected := false
450
460
prefixIdx := 0
451
461
452
462
// Selects one random spec from the range [start, end) and appends it to sampled.
453
463
collectRandomSpecFromRange := func (start , end int ) {
454
- i := start + rand .Intn (end - start )
464
+ i := start + rng .Intn (end - start )
455
465
sampled = append (sampled , specs [i ])
456
- selectedIdxs = append ( selectedIdxs , i )
466
+ selectedIndexes [ i ] = struct {}{}
457
467
}
458
468
for i , s := range specs {
459
469
if atLeastOnePerPrefix {
@@ -469,9 +479,9 @@ func selectSpecs(
469
479
}
470
480
}
471
481
472
- if rand .Float64 () < samplePct {
482
+ if rng .Float64 () < samplePct {
473
483
sampled = append (sampled , s )
474
- selectedIdxs = append ( selectedIdxs , i )
484
+ selectedIndexes [ i ] = struct {}{}
475
485
prefixSelected = true
476
486
continue
477
487
}
@@ -482,27 +492,18 @@ func selectSpecs(
482
492
}
483
493
}
484
494
485
- p := 0
486
- // The list would already be sorted were it not for the lookback to
487
- // ensure at least one test per prefix.
488
- if atLeastOnePerPrefix {
489
- sort .Ints (selectedIdxs )
490
- }
491
- // This loop depends on an ordered list as we are essentially
492
- // skipping all values in between the selected indexes.
493
- for _ , i := range selectedIdxs {
494
- for j := p ; j < i ; j ++ {
495
- s := specs [j ]
496
- if print && roachtestflags .TeamCity {
497
- fmt .Fprintf (os .Stdout , "##teamcity[testIgnored name='%s' message='excluded via sampling']\n " ,
495
+ // Print a skip message for all tests that are not selected.
496
+ for i , s := range specs {
497
+ if _ , ok := selectedIndexes [i ]; ! ok {
498
+ if stdout != nil && roachtestflags .TeamCity {
499
+ fmt .Fprintf (stdout , "##teamcity[testIgnored name='%s' message='excluded via sampling']\n " ,
498
500
s .Name )
499
501
}
500
502
501
- if print {
502
- fmt .Fprintf (os . Stdout , "--- SKIP: %s (%s)\n \t excluded via sampling\n " , s .Name , "0.00s" )
503
+ if stdout != nil {
504
+ fmt .Fprintf (stdout , "--- SKIP: %s (%s)\n \t excluded via sampling\n " , s .Name , "0.00s" )
503
505
}
504
506
}
505
- p = i + 1
506
507
}
507
508
508
509
return sampled
0 commit comments