vitess.io/vitess@v0.16.2/test.go (about)

     1  // /bin/true; exec /usr/bin/env go run "$0" "$@"
     2  
     3  /*
     4   * Copyright 2019 The Vitess Authors.
     5  
     6   * Licensed under the Apache License, Version 2.0 (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9  
    10   *     http://www.apache.org/licenses/LICENSE-2.0
    11  
    12   * Unless required by applicable law or agreed to in writing, software
    13   * distributed under the License is distributed on an "AS IS" BASIS,
    14   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    15   * See the License for the specific language governing permissions and
    16   * limitations under the License.
    17   */
    18  
    19  /*
    20  test.go is a "Go script" for running Vitess tests. It runs each test in its own
    21  Docker container for hermeticity and (potentially) parallelism. If a test fails,
    22  this script will save the output in _test/ and continue with other tests.
    23  
    24  Before using it, you should have Docker 1.5+ installed, and have your user in
    25  the group that lets you run the docker command without sudo. The first time you
    26  run against a given flavor, it may take some time for the corresponding
    27  bootstrap image (vitess/bootstrap:<flavor>) to be downloaded.
    28  
    29  It is meant to be run from the Vitess root, like so:
    30  
    31  	$ go run test.go [args]
    32  
    33  For a list of options, run:
    34  
    35  	$ go run test.go --help
    36  */
    37  package main
    38  
    39  // This Go script shouldn't rely on any packages that aren't in the standard
    40  // library, since that would require the user to bootstrap before running it.
    41  import (
    42  	"bytes"
    43  	"encoding/json"
    44  	"flag"
    45  	"fmt"
    46  	"io"
    47  	"log"
    48  	"net/http"
    49  	"net/url"
    50  	"os"
    51  	"os/exec"
    52  	"os/signal"
    53  	"path"
    54  	"path/filepath"
    55  	"sort"
    56  	"strconv"
    57  	"strings"
    58  	"sync"
    59  	"syscall"
    60  	"time"
    61  )
    62  
    63  var usage = `Usage of test.go:
    64  
    65  go run test.go [options] [test_name ...] [-- extra-py-test-args]
    66  
    67  If one or more test names are provided, run only those tests.
    68  Otherwise, run all tests in test/config.json.
    69  
    70  To pass extra args to Python tests (test/*.py), terminate the
    71  list of test names with -- and then add them at the end.
    72  
    73  For example:
    74    go run test.go test1 test2 -- --topo-flavor=etcd2
    75  `
    76  
    77  // Flags
    78  var (
    79  	flavor           = flag.String("flavor", "mysql57", "comma-separated bootstrap flavor(s) to run against (when using Docker mode). Available flavors: all,"+flavors)
    80  	bootstrapVersion = flag.String("bootstrap-version", "14.3", "the version identifier to use for the docker images")
    81  	runCount         = flag.Int("runs", 1, "run each test this many times")
    82  	retryMax         = flag.Int("retry", 3, "max number of retries, to detect flaky tests")
    83  	logPass          = flag.Bool("log-pass", false, "log test output even if it passes")
    84  	timeout          = flag.Duration("timeout", 30*time.Minute, "timeout for each test")
    85  	pull             = flag.Bool("pull", true, "re-pull the bootstrap image, in case it's been updated")
    86  	docker           = flag.Bool("docker", true, "run tests with Docker")
    87  	useDockerCache   = flag.Bool("use_docker_cache", false, "if true, create a temporary Docker image to cache the source code and the binaries generated by 'make build'. Used for execution on Travis CI.")
    88  	shard            = flag.String("shard", "", "if non-empty, run the tests whose Shard field matches value")
    89  	tag              = flag.String("tag", "", "if provided, only run tests with the given tag. Can't be combined with -shard or explicit test list")
    90  	exclude          = flag.String("exclude", "", "if provided, exclude tests containing any of the given tags (comma delimited)")
    91  	keepData         = flag.Bool("keep-data", false, "don't delete the per-test VTDATAROOT subfolders")
    92  	printLog         = flag.Bool("print-log", false, "print the log of each failed test (or all tests if -log-pass) to the console")
    93  	follow           = flag.Bool("follow", false, "print test output as it runs, instead of waiting to see if it passes or fails")
    94  	parallel         = flag.Int("parallel", 1, "number of tests to run in parallel")
    95  	skipBuild        = flag.Bool("skip-build", false, "skip running 'make build'. Assumes pre-existing binaries exist")
    96  	partialKeyspace  = flag.Bool("partial-keyspace", false, "add a second keyspace for sharded tests and mark first shard as moved to this keyspace in the shard routing rules")
    97  	// `go run test.go --dry-run --skip-build` to quickly test this file and see what tests will run
    98  	dryRun      = flag.Bool("dry-run", false, "For each test to be run, it will output the test attributes, but NOT run the tests. Useful while debugging changes to test.go (this file)")
    99  	remoteStats = flag.String("remote-stats", "", "url to send remote stats")
   100  )
   101  
   102  var (
   103  	vtDataRoot = os.Getenv("VTDATAROOT")
   104  
   105  	extraArgs []string
   106  )
   107  
   108  const (
   109  	statsFileName  = "test/stats.json"
   110  	configFileName = "test/config.json"
   111  
   112  	// List of flavors for which a bootstrap Docker image is available.
   113  	flavors = "mysql57,mysql80,percona,percona57,percona80"
   114  )
   115  
   116  // Config is the overall object serialized in test/config.json.
   117  type Config struct {
   118  	Tests map[string]*Test
   119  }
   120  
   121  // Test is an entry from the test/config.json file.
   122  type Test struct {
   123  	File          string
   124  	Args, Command []string
   125  
   126  	// Manual means it won't be run unless explicitly specified.
   127  	Manual bool
   128  
   129  	// Shard is used to split tests among workers.
   130  	Shard string
   131  
   132  	// RetryMax is the maximum number of times a test will be retried.
   133  	// If 0, flag *retryMax is used.
   134  	RetryMax int
   135  
   136  	// Tags is a list of tags that can be used to filter tests.
   137  	Tags []string
   138  
   139  	name             string
   140  	flavor           string
   141  	bootstrapVersion string
   142  	runIndex         int
   143  
   144  	pass, fail int
   145  }
   146  
   147  func (t *Test) hasTag(want string) bool {
   148  	for _, got := range t.Tags {
   149  		if got == want {
   150  			return true
   151  		}
   152  	}
   153  	return false
   154  }
   155  
   156  func (t *Test) hasAnyTag(want []string) bool {
   157  	for _, tag := range want {
   158  		if t.hasTag(tag) {
   159  			return true
   160  		}
   161  	}
   162  	return false
   163  }
   164  
   165  // run executes a single try.
   166  // dir is the location of the vitess repo to use.
   167  // dataDir is the VTDATAROOT to use for this run.
   168  // returns the combined stdout+stderr and error.
   169  func (t *Test) run(dir, dataDir string) ([]byte, error) {
   170  	if *dryRun {
   171  		fmt.Printf("Will run in dir %s(%s): %+v\n", dir, dataDir, t)
   172  		t.pass++
   173  		return nil, nil
   174  	}
   175  	testCmd := t.Command
   176  	if len(testCmd) == 0 {
   177  		if strings.Contains(fmt.Sprintf("%v", t.File), ".go") {
   178  			testCmd = []string{"tools/e2e_go_test.sh"}
   179  			testCmd = append(testCmd, t.Args...)
   180  			if *keepData {
   181  				testCmd = append(testCmd, "-keep-data")
   182  			}
   183  		} else {
   184  			testCmd = []string{"test/" + t.File, "-v", "--skip-build", "--keep-logs"}
   185  			testCmd = append(testCmd, t.Args...)
   186  		}
   187  		if *partialKeyspace {
   188  			testCmd = append(testCmd, "--partial-keyspace")
   189  		}
   190  		testCmd = append(testCmd, extraArgs...)
   191  		if *docker {
   192  			// Teardown is unnecessary since Docker kills everything.
   193  			// Go cluster doesn't recognize 'skip-teardown' flag so commenting it out for now.
   194  			// testCmd = append(testCmd, "--skip-teardown")
   195  		}
   196  	}
   197  
   198  	var cmd *exec.Cmd
   199  	if *docker {
   200  		var args []string
   201  		testArgs := strings.Join(testCmd, " ")
   202  
   203  		if *useDockerCache {
   204  			args = []string{"--use_docker_cache", cacheImage(t.flavor), t.flavor, testArgs}
   205  		} else {
   206  			// If there is no cache, we have to call 'make build' before each test.
   207  			args = []string{t.flavor, t.bootstrapVersion, "make build && " + testArgs}
   208  		}
   209  
   210  		cmd = exec.Command(path.Join(dir, "docker/test/run.sh"), args...)
   211  	} else {
   212  		cmd = exec.Command(testCmd[0], testCmd[1:]...)
   213  	}
   214  	cmd.Dir = dir
   215  
   216  	// Put everything in a unique dir, so we can copy and/or safely delete it.
   217  	// Also try to make them use different port ranges
   218  	// to mitigate failures due to zombie processes.
   219  	cmd.Env = updateEnv(os.Environ(), map[string]string{
   220  		"VTROOT":      "/vt/src/vitess.io/vitess",
   221  		"VTDATAROOT":  dataDir,
   222  		"VTPORTSTART": strconv.FormatInt(int64(getPortStart(100)), 10),
   223  	})
   224  
   225  	// Capture test output.
   226  	buf := &bytes.Buffer{}
   227  	cmd.Stdout = buf
   228  	if *follow {
   229  		cmd.Stdout = io.MultiWriter(cmd.Stdout, os.Stdout)
   230  	}
   231  	cmd.Stderr = cmd.Stdout
   232  
   233  	// Run the test.
   234  	done := make(chan error, 1)
   235  	go func() {
   236  		done <- cmd.Run()
   237  	}()
   238  
   239  	// Wait for it to finish.
   240  	var runErr error
   241  	timer := time.NewTimer(*timeout)
   242  	defer timer.Stop()
   243  	select {
   244  	case runErr = <-done:
   245  		if runErr == nil {
   246  			t.pass++
   247  		} else {
   248  			t.fail++
   249  		}
   250  	case <-timer.C:
   251  		t.logf("timeout exceeded")
   252  		cmd.Process.Signal(syscall.SIGINT)
   253  		t.fail++
   254  		runErr = <-done
   255  	}
   256  	return buf.Bytes(), runErr
   257  }
   258  
   259  func (t *Test) logf(format string, v ...any) {
   260  	if *runCount > 1 {
   261  		log.Printf("%v.%v[%v/%v]: %v", t.flavor, t.name, t.runIndex+1, *runCount, fmt.Sprintf(format, v...))
   262  	} else {
   263  		log.Printf("%v.%v: %v", t.flavor, t.name, fmt.Sprintf(format, v...))
   264  	}
   265  }
   266  
   267  func loadOneConfig(fileName string) (*Config, error) {
   268  	config2 := &Config{}
   269  	configData, err := os.ReadFile(fileName)
   270  	if err != nil {
   271  		log.Fatalf("Can't read config file %s: %v", fileName, err)
   272  		return nil, err
   273  	}
   274  	if err := json.Unmarshal(configData, config2); err != nil {
   275  		log.Fatalf("Can't parse config file: %v", err)
   276  		return nil, err
   277  	}
   278  	return config2, nil
   279  
   280  }
   281  
   282  // Get test configs.
   283  func loadConfig() (*Config, error) {
   284  	config := &Config{Tests: make(map[string]*Test)}
   285  	matches, _ := filepath.Glob("test/config*.json")
   286  	for _, configFile := range matches {
   287  		config2, err := loadOneConfig(configFile)
   288  		if err != nil {
   289  			return nil, err
   290  		}
   291  		if config2 == nil {
   292  			log.Fatalf("could not load config file: %s", configFile)
   293  		}
   294  		for key, val := range config2.Tests {
   295  			config.Tests[key] = val
   296  		}
   297  	}
   298  	return config, nil
   299  }
   300  
   301  func main() {
   302  	flag.Usage = func() {
   303  		os.Stderr.WriteString(usage)
   304  		os.Stderr.WriteString("\nOptions:\n")
   305  		flag.PrintDefaults()
   306  	}
   307  	flag.Parse()
   308  
   309  	// Sanity checks.
   310  	if *docker {
   311  		if *flavor == "all" {
   312  			*flavor = flavors
   313  		}
   314  		if *flavor == "" {
   315  			log.Fatalf("Must provide at least one -flavor when using -docker mode. Available flavors: all,%v", flavors)
   316  		}
   317  	}
   318  	if *parallel < 1 {
   319  		log.Fatalf("Invalid -parallel value: %v", *parallel)
   320  	}
   321  	if *parallel > 1 && !*docker {
   322  		log.Fatalf("Can't use -parallel value > 1 when -docker=false")
   323  	}
   324  	if *useDockerCache && !*docker {
   325  		log.Fatalf("Can't use -use_docker_cache when -docker=false")
   326  	}
   327  
   328  	startTime := time.Now()
   329  
   330  	// Make output directory.
   331  	outDir := path.Join("_test", fmt.Sprintf("%v.%v", startTime.Format("20060102-150405"), os.Getpid()))
   332  	if err := os.MkdirAll(outDir, os.FileMode(0755)); err != nil {
   333  		log.Fatalf("Can't create output directory: %v", err)
   334  	}
   335  	logFile, err := os.OpenFile(path.Join(outDir, "test.log"), os.O_RDWR|os.O_CREATE, 0644)
   336  	if err != nil {
   337  		log.Fatalf("Can't create log file: %v", err)
   338  	}
   339  	log.SetOutput(io.MultiWriter(os.Stderr, logFile))
   340  	log.Printf("Output directory: %v", outDir)
   341  
   342  	var config *Config
   343  	if config, err = loadConfig(); err != nil {
   344  		log.Fatalf("Could not load test config: %+v", err)
   345  	}
   346  
   347  	flavors := []string{"local"}
   348  
   349  	if *docker && !*dryRun {
   350  		log.Printf("Bootstrap flavor(s): %v", *flavor)
   351  
   352  		flavors = strings.Split(*flavor, ",")
   353  
   354  		// Re-pull image(s).
   355  		if *pull {
   356  			var wg sync.WaitGroup
   357  			for _, flavor := range flavors {
   358  				wg.Add(1)
   359  				go func(flavor string) {
   360  					defer wg.Done()
   361  					image := "vitess/bootstrap:" + *bootstrapVersion + "-" + flavor
   362  					pullTime := time.Now()
   363  					log.Printf("Pulling %v...", image)
   364  					cmd := exec.Command("docker", "pull", image)
   365  					if out, err := cmd.CombinedOutput(); err != nil {
   366  						log.Fatalf("Can't pull image %v: %v\n%s", image, err, out)
   367  					}
   368  					log.Printf("Image %v pulled in %v", image, round(time.Since(pullTime)))
   369  				}(flavor)
   370  			}
   371  			wg.Wait()
   372  		}
   373  	} else {
   374  		if vtDataRoot == "" {
   375  			log.Fatalf("VTDATAROOT env var must be set in -docker=false mode. Make sure to source dev.env.")
   376  		}
   377  	}
   378  
   379  	// Pick the tests to run.
   380  	var testArgs []string
   381  	testArgs, extraArgs = splitArgs(flag.Args(), "--")
   382  	tests := selectedTests(testArgs, config)
   383  
   384  	// Duplicate tests for run count.
   385  	if *runCount > 1 {
   386  		var dup []*Test
   387  		for _, t := range tests {
   388  			for i := 0; i < *runCount; i++ {
   389  				// Make a copy, since they're pointers.
   390  				test := *t
   391  				test.runIndex = i
   392  				dup = append(dup, &test)
   393  			}
   394  		}
   395  		tests = dup
   396  	}
   397  
   398  	// Duplicate tests for flavors.
   399  	var dup []*Test
   400  	for _, flavor := range flavors {
   401  		for _, t := range tests {
   402  			test := *t
   403  			test.flavor = flavor
   404  			test.bootstrapVersion = *bootstrapVersion
   405  			dup = append(dup, &test)
   406  		}
   407  	}
   408  	tests = dup
   409  
   410  	vtRoot := "."
   411  	tmpDir := ""
   412  	if *docker && !*dryRun {
   413  		// Copy working repo to tmpDir.
   414  		// This doesn't work outside Docker since it messes up GOROOT.
   415  		tmpDir, err = os.MkdirTemp(os.TempDir(), "vt_")
   416  		if err != nil {
   417  			log.Fatalf("Can't create temp dir in %v", os.TempDir())
   418  		}
   419  		log.Printf("Copying working repo to temp dir %v", tmpDir)
   420  		if out, err := exec.Command("cp", "-R", ".", tmpDir).CombinedOutput(); err != nil {
   421  			log.Fatalf("Can't copy working repo to temp dir %v: %v: %s", tmpDir, err, out)
   422  		}
   423  		// The temp copy needs permissive access so the Docker user can read it.
   424  		if out, err := exec.Command("chmod", "-R", "go=u", tmpDir).CombinedOutput(); err != nil {
   425  			log.Printf("Can't set permissions on temp dir %v: %v: %s", tmpDir, err, out)
   426  		}
   427  		vtRoot = tmpDir
   428  	} else if *skipBuild {
   429  		log.Printf("Skipping build...")
   430  	} else {
   431  		// Since we're sharing the working dir, do the build once for all tests.
   432  		log.Printf("Running make build...")
   433  		if out, err := exec.Command("make", "build").CombinedOutput(); err != nil {
   434  			log.Fatalf("make build failed: %v\n%s", err, out)
   435  		}
   436  	}
   437  
   438  	if *useDockerCache {
   439  		for _, flavor := range flavors {
   440  			start := time.Now()
   441  			log.Printf("Creating Docker cache image for flavor '%s'...", flavor)
   442  			if out, err := exec.Command("docker/test/run.sh", "--create_docker_cache", cacheImage(flavor), flavor, *bootstrapVersion, "make build").CombinedOutput(); err != nil {
   443  				log.Fatalf("Failed to create Docker cache image for flavor '%s': %v\n%s", flavor, err, out)
   444  			}
   445  			log.Printf("Creating Docker cache image took %v", round(time.Since(start)))
   446  		}
   447  	}
   448  
   449  	// Keep stats for the overall run.
   450  	var mu sync.Mutex
   451  	failed := 0
   452  	passed := 0
   453  	flaky := 0
   454  
   455  	// Listen for signals.
   456  	sigchan := make(chan os.Signal, 1)
   457  	signal.Notify(sigchan, syscall.SIGINT)
   458  
   459  	// Run tests.
   460  	stop := make(chan struct{}) // Close this to tell the runners to stop.
   461  	done := make(chan struct{}) // This gets closed when all runners have stopped.
   462  	next := make(chan *Test)    // The next test to run.
   463  	var wg sync.WaitGroup
   464  
   465  	// Send all tests into the channel.
   466  	go func() {
   467  		for _, test := range tests {
   468  			next <- test
   469  		}
   470  		close(next)
   471  	}()
   472  
   473  	// Start the requested number of parallel runners.
   474  	for i := 0; i < *parallel; i++ {
   475  		wg.Add(1)
   476  		go func() {
   477  			defer wg.Done()
   478  
   479  			for test := range next {
   480  				tryMax := *retryMax
   481  				if test.RetryMax != 0 {
   482  					tryMax = test.RetryMax
   483  				}
   484  				for try := 1; ; try++ {
   485  					select {
   486  					case <-stop:
   487  						test.logf("cancelled")
   488  						return
   489  					default:
   490  					}
   491  
   492  					if try > tryMax {
   493  						// Every try failed.
   494  						test.logf("retry limit exceeded")
   495  						mu.Lock()
   496  						failed++
   497  						mu.Unlock()
   498  						break
   499  					}
   500  
   501  					test.logf("running (try %v/%v)...", try, tryMax)
   502  
   503  					// Make a unique VTDATAROOT.
   504  					dataDir, err := os.MkdirTemp(vtDataRoot, "vt_")
   505  					if err != nil {
   506  						test.logf("Failed to create temporary subdir in VTDATAROOT: %v", vtDataRoot)
   507  						mu.Lock()
   508  						failed++
   509  						mu.Unlock()
   510  						break
   511  					}
   512  
   513  					// Run the test.
   514  					start := time.Now()
   515  					output, err := test.run(vtRoot, dataDir)
   516  					duration := time.Since(start)
   517  
   518  					// Save/print test output.
   519  					if err != nil || *logPass {
   520  						if *printLog && !*follow {
   521  							test.logf("%s\n", output)
   522  						}
   523  						outFile := fmt.Sprintf("%v.%v-%v.%v.log", test.flavor, test.name, test.runIndex+1, try)
   524  						outFilePath := path.Join(outDir, outFile)
   525  						test.logf("saving test output to %v", outFilePath)
   526  						if fileErr := os.WriteFile(outFilePath, output, os.FileMode(0644)); fileErr != nil {
   527  							test.logf("WriteFile error: %v", fileErr)
   528  						}
   529  					}
   530  
   531  					// Clean up the unique VTDATAROOT.
   532  					if !*keepData {
   533  						if err := os.RemoveAll(dataDir); err != nil {
   534  							test.logf("WARNING: can't remove temporary VTDATAROOT: %v", err)
   535  						}
   536  					}
   537  
   538  					if err != nil {
   539  						// This try failed.
   540  						test.logf("FAILED (try %v/%v) in %v: %v", try, tryMax, round(duration), err)
   541  						mu.Lock()
   542  						testFailed(test.name)
   543  						mu.Unlock()
   544  						continue
   545  					}
   546  
   547  					mu.Lock()
   548  					testPassed(test.name, duration)
   549  
   550  					if try == 1 {
   551  						// Passed on the first try.
   552  						test.logf("PASSED in %v", round(duration))
   553  						passed++
   554  					} else {
   555  						// Passed, but not on the first try.
   556  						test.logf("FLAKY (1/%v passed in %v)", try, round(duration))
   557  						flaky++
   558  						testFlaked(test.name, try)
   559  					}
   560  					mu.Unlock()
   561  					break
   562  				}
   563  			}
   564  		}()
   565  	}
   566  
   567  	// Close the done channel when all the runners stop.
   568  	// This lets us select on wg.Wait().
   569  	go func() {
   570  		wg.Wait()
   571  		close(done)
   572  	}()
   573  
   574  	// Stop the loop and kill child processes if we get a signal.
   575  	select {
   576  	case <-sigchan:
   577  		log.Printf("interrupted: skip remaining tests and wait for current test to tear down")
   578  		signal.Stop(sigchan)
   579  		// Stop the test loop and wait for it to exit.
   580  		// Running tests already get the SIGINT themselves.
   581  		// We mustn't send it again, or it'll abort the teardown process too early.
   582  		close(stop)
   583  		<-done
   584  	case <-done:
   585  	}
   586  
   587  	// Clean up temp dir.
   588  	if tmpDir != "" {
   589  		log.Printf("Removing temp dir %v", tmpDir)
   590  		if err := os.RemoveAll(tmpDir); err != nil {
   591  			log.Printf("Failed to remove temp dir: %v", err)
   592  		}
   593  	}
   594  	// Remove temporary Docker cache image.
   595  	if *useDockerCache {
   596  		for _, flavor := range flavors {
   597  			log.Printf("Removing temporary Docker cache image for flavor '%s'", flavor)
   598  			if out, err := exec.Command("docker", "rmi", cacheImage(flavor)).CombinedOutput(); err != nil {
   599  				log.Printf("WARNING: Failed to delete Docker cache image: %v\n%s", err, out)
   600  			}
   601  		}
   602  	}
   603  
   604  	// Print summary.
   605  	log.Print(strings.Repeat("=", 60))
   606  	for _, t := range tests {
   607  		tname := t.flavor + "." + t.name
   608  		switch {
   609  		case t.pass > 0 && t.fail == 0:
   610  			log.Printf("%-40s\tPASS", tname)
   611  		case t.pass > 0 && t.fail > 0:
   612  			log.Printf("%-40s\tFLAKY (%v/%v failed)", tname, t.fail, t.pass+t.fail)
   613  		case t.pass == 0 && t.fail > 0:
   614  			log.Printf("%-40s\tFAIL (%v tries)", tname, t.fail)
   615  		case t.pass == 0 && t.fail == 0:
   616  			log.Printf("%-40s\tSKIPPED", tname)
   617  		}
   618  	}
   619  	log.Print(strings.Repeat("=", 60))
   620  	skipped := len(tests) - passed - flaky - failed
   621  	log.Printf("%v PASSED, %v FLAKY, %v FAILED, %v SKIPPED", passed, flaky, failed, skipped)
   622  	log.Printf("Total time: %v", round(time.Since(startTime)))
   623  
   624  	if failed > 0 || skipped > 0 {
   625  		os.Exit(1)
   626  	}
   627  }
   628  
   629  func updateEnv(orig []string, updates map[string]string) []string {
   630  	var env []string
   631  	for _, v := range orig {
   632  		parts := strings.SplitN(v, "=", 2)
   633  		if _, ok := updates[parts[0]]; !ok {
   634  			env = append(env, v)
   635  		}
   636  	}
   637  	for k, v := range updates {
   638  		env = append(env, k+"="+v)
   639  	}
   640  	return env
   641  }
   642  
   643  // cacheImage returns the flavor-specific name of the Docker cache image.
   644  func cacheImage(flavor string) string {
   645  	return fmt.Sprintf("vitess/bootstrap:rm_%s_test_cache_do_NOT_push", flavor)
   646  }
   647  
   648  type Stats struct {
   649  	TestStats map[string]TestStats
   650  }
   651  
   652  type TestStats struct {
   653  	Pass, Fail, Flake int
   654  	PassTime          time.Duration
   655  
   656  	name string
   657  }
   658  
   659  func sendStats(values url.Values) {
   660  	if *remoteStats != "" {
   661  		log.Printf("Sending remote stats to %v", *remoteStats)
   662  		resp, err := http.PostForm(*remoteStats, values)
   663  		if err != nil {
   664  			log.Printf("Can't send remote stats: %v", err)
   665  		}
   666  		defer resp.Body.Close()
   667  	}
   668  }
   669  
   670  func testPassed(name string, passTime time.Duration) {
   671  	sendStats(url.Values{
   672  		"test":     {name},
   673  		"result":   {"pass"},
   674  		"duration": {passTime.String()},
   675  	})
   676  	updateTestStats(name, func(ts *TestStats) {
   677  		totalTime := int64(ts.PassTime)*int64(ts.Pass) + int64(passTime)
   678  		ts.Pass++
   679  		ts.PassTime = time.Duration(totalTime / int64(ts.Pass))
   680  	})
   681  }
   682  
   683  func testFailed(name string) {
   684  	sendStats(url.Values{
   685  		"test":   {name},
   686  		"result": {"fail"},
   687  	})
   688  	updateTestStats(name, func(ts *TestStats) {
   689  		ts.Fail++
   690  	})
   691  }
   692  
   693  func testFlaked(name string, try int) {
   694  	sendStats(url.Values{
   695  		"test":   {name},
   696  		"result": {"flake"},
   697  		"try":    {strconv.FormatInt(int64(try), 10)},
   698  	})
   699  	updateTestStats(name, func(ts *TestStats) {
   700  		ts.Flake += try - 1
   701  	})
   702  }
   703  
   704  func updateTestStats(name string, update func(*TestStats)) {
   705  	var stats Stats
   706  
   707  	data, err := os.ReadFile(statsFileName)
   708  	if err != nil {
   709  		log.Print("Can't read stats file, starting new one.")
   710  	} else {
   711  		if err := json.Unmarshal(data, &stats); err != nil {
   712  			log.Printf("Can't parse stats file: %v", err)
   713  			return
   714  		}
   715  	}
   716  
   717  	if stats.TestStats == nil {
   718  		stats.TestStats = make(map[string]TestStats)
   719  	}
   720  	ts := stats.TestStats[name]
   721  	update(&ts)
   722  	stats.TestStats[name] = ts
   723  
   724  	data, err = json.MarshalIndent(stats, "", "\t")
   725  	if err != nil {
   726  		log.Printf("Can't encode stats file: %v", err)
   727  		return
   728  	}
   729  	if err := os.WriteFile(statsFileName, data, 0644); err != nil {
   730  		log.Printf("Can't write stats file: %v", err)
   731  	}
   732  }
   733  
   734  type ByPassTime []TestStats
   735  
   736  func (a ByPassTime) Len() int           { return len(a) }
   737  func (a ByPassTime) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
   738  func (a ByPassTime) Less(i, j int) bool { return a[i].PassTime > a[j].PassTime }
   739  
   740  func getTestsSorted(names []string, testMap map[string]*Test) []*Test {
   741  	sort.Strings(names)
   742  	var tests []*Test
   743  	for _, name := range names {
   744  		t := testMap[name]
   745  		t.name = name
   746  		tests = append(tests, t)
   747  	}
   748  	return tests
   749  }
   750  
   751  func selectedTests(args []string, config *Config) []*Test {
   752  	var tests []*Test
   753  	excludedTests := strings.Split(*exclude, ",")
   754  	if *shard != "" {
   755  		// Run the tests in a given shard.
   756  		// This can be combined with positional args.
   757  		var names []string
   758  		for name, t := range config.Tests {
   759  			if t.Shard == *shard && !t.Manual && (*exclude == "" || !t.hasAnyTag(excludedTests)) {
   760  				t.name = name
   761  				names = append(names, name)
   762  			}
   763  		}
   764  		tests = getTestsSorted(names, config.Tests)
   765  	}
   766  	if len(args) > 0 {
   767  		// Positional args for manual selection.
   768  		for _, name := range args {
   769  			t, ok := config.Tests[name]
   770  			if !ok {
   771  				tests := make([]string, len(config.Tests))
   772  
   773  				i := 0
   774  				for k := range config.Tests {
   775  					tests[i] = k
   776  					i++
   777  				}
   778  
   779  				sort.Strings(tests)
   780  
   781  				log.Fatalf("Unknown test: %v\nAvailable tests are: %v", name, strings.Join(tests, ", "))
   782  			}
   783  			t.name = name
   784  			tests = append(tests, t)
   785  		}
   786  	}
   787  	if len(args) == 0 && *shard == "" {
   788  		// Run all tests.
   789  		var names []string
   790  		for name, t := range config.Tests {
   791  			if !t.Manual && (*tag == "" || t.hasTag(*tag)) && (*exclude == "" || !t.hasAnyTag(excludedTests)) {
   792  				names = append(names, name)
   793  			}
   794  		}
   795  		tests = getTestsSorted(names, config.Tests)
   796  	}
   797  	return tests
   798  }
   799  
   800  var (
   801  	port      = 16000
   802  	portMutex sync.Mutex
   803  )
   804  
   805  func getPortStart(size int) int {
   806  	portMutex.Lock()
   807  	defer portMutex.Unlock()
   808  
   809  	start := port
   810  	port += size
   811  	return start
   812  }
   813  
   814  // splitArgs splits a list of args at the first appearance of tok.
   815  func splitArgs(all []string, tok string) (args, extraArgs []string) {
   816  	extra := false
   817  	for _, arg := range all {
   818  		if extra {
   819  			extraArgs = append(extraArgs, arg)
   820  			continue
   821  		}
   822  		if arg == tok {
   823  			extra = true
   824  			continue
   825  		}
   826  		args = append(args, arg)
   827  	}
   828  	return
   829  }
   830  
   831  func round(d time.Duration) time.Duration {
   832  	return d.Round(100 * time.Millisecond)
   833  }