go.fuchsia.dev/infra@v0.0.0-20240507153436-9b593402251b/cmd/update_test_durations/update.go

// Copyright 2022 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// This tool queries ResultDB data from BigQuery to retrieve statistics about to
// the durations of Fuchsia tests. The resulting data will be uploaded to CIPD
// for use by the testsharder tool:
// https://fuchsia.googlesource.com/fuchsia/+/HEAD/tools/integration/testsharder

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/maruel/subcommands"
	"go.chromium.org/luci/auth"
	"go.fuchsia.dev/infra/functools"
	"golang.org/x/exp/maps"
)

var errZeroTotalRuns = errors.New("cannot calculate average duration for tests with zero total runs")

const (
	// For each set of duration files (public and internal), the statistics for
	// all tests will be aggregated to create a new file containing each test's
	// average duration across all builders. This file will be used by any
	// builders that don't yet have a corresponding duration file.
	defaultBuilderName = "default"

	// Every duration file will have a default entry that's applied to any test
	// that doesn't yet have an entry (probably because it is new or renamed). It
	// will have this value in its "name" field. We chose "*" to strike a balance
	// between distinguishing this entry from actual tests (text "default" would
	// be easier to mistake for an actual test) and being obviously intentional
	// and clear about its purpose (an empty string would less clearly be a
	// fallback, and might even suggest a bug in the updater).
	defaultTestName = "*"
)

func cmdRun(authOpts auth.Options) *subcommands.Command {
	return &subcommands.Command{
		UsageLine: "run -dir DIR -project PROJECT",
		ShortDesc: "write updated test duration data to a directory",
		LongDesc:  "write updated test duration data to a directory",
		CommandRun: func() subcommands.CommandRun {
			c := &runCmd{}
			c.Init(authOpts)
			return c
		},
	}
}

type runCmd struct {
	commonFlags

	previousVersionDir string
	outputDir          string
	luciProject        string
	dataWindowDays     int
}

// test is a pairing of a test and builder, along with the number of recorded
// runs and median duration within the time window. The order of the fields here
// must be kept in sync with the order of the rows returned by the BigQuery
// query.
type test struct {
	Name string `json:"name"`

	// The number of test runs included in calculating the duration.
	Runs int64 `json:"runs"`

	// The median duration of the test, in milliseconds, across all included runs.
	MedianDurationMS int64 `json:"median_duration_ms"`

	// The builder that ran this test. Test durations are separated into files
	// by builder, so it's not necessary to include the builder name in the
	// output JSON.
	Builder string `json:"-"`
}

// testDurationsMap maps from the name of a builder to a list of tests that were
// run by that builder. Each entry in the map corresponds to one file in the
// resulting test duration files.
type testDurationMap map[string][]test

// durationFileOptions contains feature flags for splitTestsByBuilder. This is
// primarily for testing purposes; turning off some features makes it easier to
// construct expected data for tests where we're making assertions that aren't
// related to those features.
//
// All features should always be enabled in production.
type durationFileOptions struct {
	includeDefaultTests   bool
	includeDefaultBuilder bool
}

func (c *runCmd) Init(defaultAuthOpts auth.Options) {
	c.commonFlags.Init(defaultAuthOpts)
	c.Flags.StringVar(
		&c.previousVersionDir,
		"previous-version-dir",
		"",
		"Directory containing the previous version of the durations files, which "+
			"will be merged with the updated durations.")
	c.Flags.StringVar(
		&c.outputDir,
		"output-dir",
		"",
		"Directory into which final duration files should be written.")
	c.Flags.StringVar(
		&c.luciProject,
		"project",
		"fuchsia",
		"LUCI project to query test durations for.")
	c.Flags.IntVar(
		&c.dataWindowDays,
		"days",
		3,
		"LUCI project to query test durations for.")
}

func (c *runCmd) parseArgs([]string) error {
	return c.commonFlags.Parse()
}

func (c *runCmd) Run(a subcommands.Application, args []string, _ subcommands.Env) int {
	if err := c.parseArgs(args); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}

	if err := c.main(context.Background()); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	return 0
}

// updateTestDurations fetches the latest test durations from BigQuery and
// uploads them to CIPD.
func (c *runCmd) main(ctx context.Context) error {
	if c.previousVersionDir == "" {
		return errors.New("flag -previous-version-dir is required")
	}
	if c.outputDir == "" {
		return errors.New("flag -output-dir is required")
	}
	fetchCtx, cancel := context.WithTimeout(ctx, queryTimeout)
	defer cancel()
	rows, err := queryLatestTestDurations(fetchCtx, c.luciProject, c.dataWindowDays)
	if err != nil {
		return err
	}

	durations, err := splitTestsByBuilder(rows, durationFileOptions{
		includeDefaultTests:   true,
		includeDefaultBuilder: true,
	})
	if err != nil {
		return err
	}

	newFileContents, err := marshalDurations(durations)
	if err != nil {
		return err
	}

	oldFileContents, err := previousVersionContents(c.previousVersionDir)
	if err != nil {
		return err
	}

	finalContents := overlayFileContents(oldFileContents, newFileContents)
	for name, contents := range finalContents {
		path := filepath.Join(c.outputDir, name)
		if err := os.WriteFile(path, contents, 0o600); err != nil {
			return fmt.Errorf("failed to write: %w", err)
		}
	}

	return nil
}

// splitTestsByBuilder takes a collection of tests and arranges them into a
// mapping corresponding to the data that should be written to each per-builder
// test duration file.
func splitTestsByBuilder(tests []test, opts durationFileOptions) (testDurationMap, error) {
	// Mapping from builder name to list of tests.
	durations := make(testDurationMap)

	if len(tests) == 0 {
		return durations, nil
	}

	// Mapping from test name to list of tests, across all builders. Only
	// used as intermediate storage for producing the default test duration file.
	testsByName := make(map[string][]test)
	for _, t := range tests {
		if t.Builder == "" {
			return nil, fmt.Errorf("test has empty builder: %+v", t)
		}
		durations[t.Builder] = append(durations[t.Builder], t)
		testsByName[t.Name] = append(testsByName[t.Name], t)
	}

	if opts.includeDefaultBuilder {
		defaultDurations, err := calculateDefaultDurations(testsByName)
		if err != nil {
			return nil, err
		}
		durations[defaultBuilderName] = defaultDurations
	}

	if opts.includeDefaultTests {
		if err := addDefaultEntries(durations); err != nil {
			return nil, err
		}
	}

	return durations, nil
}

// calculateDefaultDurations aggregates durations for tests across all builders.
// testsharder will use this default file for any unknown builder that doesn't
// yet have its own durations file.
//
// We use the average duration of all tests, rather than the median, so that
// when using this default file, the sum of all tests' expected durations is
// close to the sum of actual durations. If we used the median of all durations,
// the sum of expected durations would likely be far too low because it wouldn't
// account for the long tail of slower tests, and we would end up producing too
// few shards when executing testsharder with a target per-shard duration.
//
// See unit tests for examples.
func calculateDefaultDurations(testsByName map[string][]test) ([]test, error) {
	var defaultDurations []test
	for name, sameNameTests := range testsByName {
		var totalRuns int64
		for _, t := range sameNameTests {
			totalRuns += t.Runs
		}

		duration, err := averageDurationMS(sameNameTests)
		if err != nil {
			return nil, err
		}

		defaultDurations = append(defaultDurations, test{
			Name:             name,
			Builder:          defaultBuilderName,
			Runs:             totalRuns,
			MedianDurationMS: duration,
		})
	}

	// The SQL query ensures that all other builders' tests are sorted, so we
	// sort these too for consistency.
	sort.Slice(defaultDurations, func(i, j int) bool {
		return defaultDurations[i].Name < defaultDurations[j].Name
	})

	return defaultDurations, nil
}

// addDefaultEntries adds an new entry for each builder at index zero that
// will be applied by testsharder to any test that doesn't have its own entry,
// probably because it's a new test or has been renamed.
//
// The average of all existing tests' median durations is probably a good
// estimate of the duration for any such tests. Average of medians is better
// than median of medians in the case where many new tests are added, assuming
// that the distribution of durations of the new tests is similar to the
// distribution for existing tests. (If only a couple tests are added, it
// doesn't make much of a difference either way). The median of medians might be
// closer to the actual duration of *most* of the new tests, but it doesn't take
// into account the fact that a few of the new tests are probably much longer
// than the rest. This would lead to all the new tests being put into the same
// shard, which would always time out because it has too many long tests.
func addDefaultEntries(durations testDurationMap) error {
	for builder, builderTests := range durations {
		defaultDuration, err := averageDurationMS(builderTests)
		if err != nil {
			return err
		}
		defaultTestEntry := test{
			Name:             defaultTestName,
			Builder:          builder,
			MedianDurationMS: defaultDuration,
		}
		durations[builder] = append([]test{defaultTestEntry}, builderTests...)
	}
	return nil
}

// averageDurationMS calculates the average of median durations, weighted by
// runs, for all the given tests.
func averageDurationMS(tests []test) (int64, error) {
	var totalDurationMS, totalRuns int64
	for _, t := range tests {
		totalRuns += t.Runs
		totalDurationMS += t.MedianDurationMS * t.Runs
	}
	if totalRuns == 0 {
		// This likely indicates a bug somewhere in the SQL query or Go code.
		return 0, fmt.Errorf("%w: %+v", errZeroTotalRuns, tests)
	}
	avg := float64(totalDurationMS) / float64(totalRuns)
	return int64(math.Round(avg)), nil
}

func marshalDurations(durations testDurationMap) (map[string][]byte, error) {
	files := make(map[string][]byte, len(durations))
	for builder, tests := range durations {
		functools.SortBy(tests, func(t test) string {
			return t.Name
		})
		b, err := json.Marshal(tests)
		if err != nil {
			return nil, err
		}
		fileName := fmt.Sprintf("%s.json", builder)
		files[fileName] = b
	}

	return files, nil
}

// previousVersionContents downloads the version of `pkg` identified by `ref`
// and returns a mapping from file basename to file contents for all the non-
// hidden files in the root of the package.
func previousVersionContents(previousDurationsDir string) (map[string][]byte, error) {
	files, err := os.ReadDir(previousDurationsDir)
	if err != nil {
		return nil, err
	}

	contents := make(map[string][]byte)
	for _, file := range files {
		name := file.Name()
		// Ignore hidden files and directories to ensure we only copy duration
		// files into the updated package, and not any special files installed
		// by CIPD itself.
		if strings.HasPrefix(name, ".") || file.IsDir() {
			continue
		}
		b, err := os.ReadFile(filepath.Join(previousDurationsDir, name))
		if err != nil {
			return nil, err
		}
		contents[name] = b
	}

	return contents, nil
}

// overlayFileContents takes an in-memory listing of the current contents of a
// directory and overlays that listing with a listing of new files. It leaves in
// place any entries from `oldFiles` that do not have a corresponding entry in
// `newFiles`.
//
// The purpose of this is to keep around duration files from previous package
// versions even if we don't have any data for their builders from this run of
// the updater. By keeping the old files around, we'll ensure that future builds
// still have test duration data to use.
//
// This is useful if, for example, a builder is paused or breaks for a few days
// and we stop getting data from it, but then it gets fixed/restored. We still
// want the builder to have duration data to use when it comes back online, so
// we intentionally keep its duration file around.
//
// TODO(olivernewman): Set up a garbage collection system so we don't keep old
// files around forever.
func overlayFileContents(oldFiles, newFiles map[string][]byte) map[string][]byte {
	result := make(map[string][]byte)
	maps.Copy(result, oldFiles)
	maps.Copy(result, newFiles)
	return result
}