cmd/testwrapper, tstest: move test sharding out of test code

Previously, sharding required tests to opt in by calling tstest.Shard,
which used a process-global counter to assign each test to a shard.
This had two problems: most tests didn't call it, so they ran on every
shard (defeating the purpose), and shard assignments were unstable
(depended on call order, so adding a test could reshuffle others).

Remove tstest.Shard and tstest.SkipOnUnshardedCI entirely. Instead,
have testwrapper implement sharding automatically for all tests: when
TS_TEST_SHARD=N/M is set, it uses "go list -json" (no compilation) to
find test source files, scans them for top-level Test/Benchmark/
Example/Fuzz function names, and filters by fnv32a(name) % M == N-1.
The filtered names are passed as an anchored -run regex to go test.

Using go list instead of "go test -list" avoids linking the test binary
twice (Go's build cache does not cache test binary linking).

Fixes #19886

Change-Id: I62ab7b3d757324d4c5fd0b5de50c1e3742681791
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>

cmd/testwrapper/testwrapper.go

@@ -15,12 +15,16 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"hash/fnv"
 	"io"
 	"log"
 	"os"
 	"os/exec"
+	"path/filepath"
+	"regexp"
 	"slices"
 	"sort"
+	"strconv"
 	"strings"
 	"time"
 
@@ -83,6 +87,94 @@ type goTestOutput struct {
 
 var debug = os.Getenv("TS_TESTWRAPPER_DEBUG") != ""
 
+// testsForShard returns the test names in pkg that belong to the given shard
+// spec (e.g. "2/3"). It uses "go list -json" to find test source files (no
+// compilation) and scans them for top-level test function names, assigning
+// each to a shard by hashing. Returns nil if the spec is invalid or if
+// listing fails (the main run will surface the error).
+func testsForShard(ctx context.Context, pkg, shardSpec string) ([]string, error) {
+	a, b, ok := strings.Cut(shardSpec, "/")
+	if !ok {
+		return nil, nil
+	}
+	wantShard, err := strconv.Atoi(a)
+	if err != nil || wantShard < 1 {
+		return nil, nil
+	}
+	shards, err := strconv.Atoi(b)
+	if err != nil || shards < 1 {
+		return nil, nil
+	}
+
+	out, err := exec.CommandContext(ctx, "go", "list", "-json", pkg).Output()
+	if err != nil {
+		// Errors will be surfaced by the main test run.
+		return nil, nil
+	}
+
+	type pkgJSON struct {
+		Dir          string
+		TestGoFiles  []string
+		XTestGoFiles []string
+	}
+
+	seen := map[string]bool{}
+	var result []string
+
+	dec := json.NewDecoder(bytes.NewReader(out))
+	for dec.More() {
+		var p pkgJSON
+		if err := dec.Decode(&p); err != nil {
+			break
+		}
+		for _, f := range append(p.TestGoFiles, p.XTestGoFiles...) {
+			names, err := testFuncNames(filepath.Join(p.Dir, f))
+			if err != nil {
+				continue
+			}
+			for _, name := range names {
+				if seen[name] {
+					continue
+				}
+				seen[name] = true
+				h := fnv.New32a()
+				io.WriteString(h, name)
+				if int(h.Sum32()%uint32(shards)) == wantShard-1 {
+					result = append(result, name)
+				}
+			}
+		}
+	}
+	return result, nil
+}
+
+// testFuncNames scans a Go source file and returns the names of all top-level
+// test functions (Test*, Benchmark*, Example*, Fuzz*).
+func testFuncNames(path string) ([]string, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+	var names []string
+	sc := bufio.NewScanner(f)
+	for sc.Scan() {
+		rest, ok := strings.CutPrefix(sc.Text(), "func ")
+		if !ok {
+			continue
+		}
+		for _, prefix := range []string{"Test", "Benchmark", "Example", "Fuzz"} {
+			if strings.HasPrefix(rest, prefix) {
+				if i := strings.IndexByte(rest, '('); i > 0 {
+					names = append(names, rest[:i])
+				}
+				break
+			}
+		}
+	}
+	return names, sc.Err()
+}
+
 // runTests runs the tests in pt and sends the results on ch. It sends a
 // testAttempt for each test and a final testAttempt per pkg with pkgFinished
 // set to true. Package build errors will not emit a testAttempt (as no valid
@@ -94,8 +186,24 @@ func runTests(ctx context.Context, attempt int, pt *packageTests, goTestArgs, te
 	args = append(args, goTestArgs...)
 	args = append(args, pt.Pattern)
 	if len(pt.Tests) > 0 {
+		// Specific tests requested (e.g. flaky test retry).
 		runArg := strings.Join(pt.Tests, "|")
 		args = append(args, "--run", runArg)
+	} else if shardSpec := os.Getenv("TS_TEST_SHARD"); shardSpec != "" {
+		// Automatic test-name sharding: list tests and filter by hash.
+		shardTests, err := testsForShard(ctx, pt.Pattern, shardSpec)
+		if err != nil {
+			return err
+		}
+		if len(shardTests) == 0 {
+			ch <- &testAttempt{pkg: pt.Pattern, outcome: "skip", pkgFinished: true}
+			return nil
+		}
+		quoted := make([]string, len(shardTests))
+		for i, name := range shardTests {
+			quoted[i] = regexp.QuoteMeta(name)
+		}
+		args = append(args, "--run", "^("+strings.Join(quoted, "|")+")$")
 	}
 	args = append(args, testArgs...)
 	args = append(args, "-json")
@@ -103,9 +211,6 @@ func runTests(ctx context.Context, attempt int, pt *packageTests, goTestArgs, te
 		fmt.Println("running", strings.Join(args, " "))
 	}
 	cmd := exec.CommandContext(ctx, "go", args...)
-	if len(pt.Tests) > 0 {
-		cmd.Env = append(os.Environ(), "TS_TEST_SHARD=") // clear test shard; run all tests we say to run
-	}
 	r, err := cmd.StdoutPipe()
 	if err != nil {
 		log.Printf("error creating stdout pipe: %v", err)
@@ -113,7 +218,9 @@ func runTests(ctx context.Context, attempt int, pt *packageTests, goTestArgs, te
 	defer r.Close()
 	cmd.Stderr = os.Stderr
 
-	cmd.Env = os.Environ()
+	cmd.Env = slices.DeleteFunc(os.Environ(), func(s string) bool {
+		return strings.HasPrefix(s, "TS_TEST_SHARD=")
+	})
 	cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%d", flakytest.FlakeAttemptEnv, attempt))
 
 	if err := cmd.Start(); err != nil {