roachtest: ensure all skipped tests are logged

Previously, there was a bug where the last tests after the selected indexes were
not being logged as skipped. This change adds the length of the specs to the
list of selected indexes to ensure that the loop will skip all the tests between
the selected indexes, including the tail end of the tests.

Epic: None
Release note: None

Author: herkolategan

pkg/cmd/roachtest/main.go

@@ -8,12 +8,12 @@ package main
 import (
 	"context"
 	"fmt"
+	"io"
 	"math"
 	"math/rand"
 	"os"
 	"os/user"
 	"regexp"
-	"sort"
 	"strings"
 	"time"
 
@@ -25,6 +25,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/testselector"
 	"github.com/cockroachdb/cockroach/pkg/roachprod"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/config"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/errors"
 	_ "github.com/lib/pq" // register postgres driver
 	"github.com/spf13/cobra"
@@ -315,7 +316,12 @@ func testsToRun(
 		}
 	}
 
-	return selectSpecs(notSkipped, selectProbability, true, print), nil
+	var stdout io.Writer
+	if print {
+		stdout = os.Stdout
+	}
+	rng, _ := randutil.NewPseudoRand()
+	return selectSpecs(notSkipped, rng, selectProbability, true, stdout), nil
 }
 
 // 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
 // testRegistryImpl.AllTests().
 // TODO(smg260): Perhaps expose `atLeastOnePerPrefix` via CLI
 func selectSpecs(
-	specs []registry.TestSpec, samplePct float64, atLeastOnePerPrefix bool, print bool,
+	specs []registry.TestSpec,
+	rng *rand.Rand,
+	samplePct float64,
+	atLeastOnePerPrefix bool,
+	stdout io.Writer,
 ) []registry.TestSpec {
 	if samplePct == 1 || len(specs) == 0 {
 		return specs
 	}
 
 	var sampled []registry.TestSpec
-	var selectedIdxs []int
+	selectedIndexes := make(map[int]struct{})
 
 	prefix := strings.Split(specs[0].Name, "/")[0]
 	prefixSelected := false
 	prefixIdx := 0
 
 	// Selects one random spec from the range [start, end) and appends it to sampled.
 	collectRandomSpecFromRange := func(start, end int) {
-		i := start + rand.Intn(end-start)
+		i := start + rng.Intn(end-start)
 		sampled = append(sampled, specs[i])
-		selectedIdxs = append(selectedIdxs, i)
+		selectedIndexes[i] = struct{}{}
 	}
 	for i, s := range specs {
 		if atLeastOnePerPrefix {
@@ -469,9 +479,9 @@ func selectSpecs(
 			}
 		}
 
-		if rand.Float64() < samplePct {
+		if rng.Float64() < samplePct {
 			sampled = append(sampled, s)
-			selectedIdxs = append(selectedIdxs, i)
+			selectedIndexes[i] = struct{}{}
 			prefixSelected = true
 			continue
 		}
@@ -482,27 +492,18 @@ func selectSpecs(
 		}
 	}
 
-	p := 0
-	// The list would already be sorted were it not for the lookback to
-	// ensure at least one test per prefix.
-	if atLeastOnePerPrefix {
-		sort.Ints(selectedIdxs)
-	}
-	// This loop depends on an ordered list as we are essentially
-	// skipping all values in between the selected indexes.
-	for _, i := range selectedIdxs {
-		for j := p; j < i; j++ {
-			s := specs[j]
-			if print && roachtestflags.TeamCity {
-				fmt.Fprintf(os.Stdout, "##teamcity[testIgnored name='%s' message='excluded via sampling']\n",
+	// Print a skip message for all tests that are not selected.
+	for i, s := range specs {
+		if _, ok := selectedIndexes[i]; !ok {
+			if stdout != nil && roachtestflags.TeamCity {
+				fmt.Fprintf(stdout, "##teamcity[testIgnored name='%s' message='excluded via sampling']\n",
 					s.Name)
 			}
 
-			if print {
-				fmt.Fprintf(os.Stdout, "--- SKIP: %s (%s)\n\texcluded via sampling\n", s.Name, "0.00s")
+			if stdout != nil {
+				fmt.Fprintf(stdout, "--- SKIP: %s (%s)\n\texcluded via sampling\n", s.Name, "0.00s")
 			}
 		}
-		p = i + 1
 	}
 
 	return sampled