tools/integration/testsharder/cmd/main.go

// Copyright 2018 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.

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	"go.fuchsia.dev/fuchsia/tools/build"
	"go.fuchsia.dev/fuchsia/tools/integration/testsharder"
	"go.fuchsia.dev/fuchsia/tools/lib/color"
	"go.fuchsia.dev/fuchsia/tools/lib/flagmisc"
	"go.fuchsia.dev/fuchsia/tools/lib/logger"
)

var (
	buildDir                       string
	outputFile                     string
	mode                           testsharder.Mode = testsharder.Normal
	tags                           flagmisc.StringsValue
	modifiersPath                  string
	targetTestCount                int
	targetDurationSecs             int
	perTestTimeoutSecs             int
	maxShardsPerEnvironment        int
	affectedTestsPath              string
	affectedTestsMaxAttempts       int
	affectedTestsMultiplyThreshold int
	affectedOnly                   bool
	realmLabel                     string
	hermeticDeps                   bool
	pave                           bool
)

func usage() {
	fmt.Printf(`testsharder [flags]

Shards tests produced by a build.
For more information on the modes in which the testsharder may be run, see
https://pkg.go.dev/go.fuchsia.dev/fuchsia/tools/integration/testsharder#Mode.
`)
}

func init() {
	flag.StringVar(&buildDir, "build-dir", "", "path to the fuchsia build directory root (required)")
	flag.StringVar(&outputFile, "output-file", "", "path to a file which will contain the shards as JSON, default is stdout")
	flag.Var(&mode, "mode", "mode in which to run the testsharder (e.g., normal or restricted).")
	flag.Var(&tags, "tag", "environment tags on which to filter; only the tests that match all tags will be sharded")
	flag.StringVar(&modifiersPath, "modifiers", "", "path to the json manifest containing tests to modify")
	flag.IntVar(&targetDurationSecs, "target-duration-secs", 0, "approximate duration that each shard should run in")
	flag.IntVar(&maxShardsPerEnvironment, "max-shards-per-env", 8, "maximum shards allowed per environment. If <= 0, no max will be set")
	// TODO(fxbug.dev/10456): Support different timeouts for different tests.
	flag.IntVar(&perTestTimeoutSecs, "per-test-timeout-secs", 0, "per-test timeout, applied to all tests. If <= 0, no timeout will be set")
	// Despite being a misnomer, this argument is still called -max-shard-size
	// for legacy reasons. If it becomes confusing, we can create a new
	// target_test_count fuchsia.proto field and do a soft transition with the
	// recipes to start setting the renamed argument instead.
	flag.IntVar(&targetTestCount, "max-shard-size", 0, "target number of tests per shard. If <= 0, will be ignored. Otherwise, tests will be placed into more, smaller shards")
	flag.StringVar(&affectedTestsPath, "affected-tests", "", "path to a file containing names of tests affected by the change being tested. One test name per line.")
	flag.IntVar(&affectedTestsMaxAttempts, "affected-tests-max-attempts", 2, "maximum attempts for each affected test. Only applied to tests that are not multiplied")
	flag.IntVar(&affectedTestsMultiplyThreshold, "affected-tests-multiply-threshold", 0, "if there are <= this many tests in -affected-tests, they may be multplied "+
		"(modified to run many times in a separate shard), but only be multiplied if allowed by certain constraints designed to minimize false rejections and bot demand.")
	flag.BoolVar(&affectedOnly, "affected-only", false, "whether to create test shards for only the affected tests found in either the modifiers file or the affected-tests file.")
	flag.StringVar(&realmLabel, "realm-label", "", "applies this realm label to the output sharded json file generated by testsharder. If empty, no realm label is applied.")
	flag.BoolVar(&hermeticDeps, "hermetic-deps", false, "whether to add all the images and blobs used by the shard as dependencies")
	flag.BoolVar(&pave, "pave", false, "whether the shards generated should pave or netboot fuchsia")
	flag.Usage = usage
}

func main() {
	flag.Parse()
	l := logger.NewLogger(logger.ErrorLevel, color.NewColor(color.ColorAuto), os.Stdout, os.Stderr, "")
	// testsharder is expected to complete quite quickly, so it's generally not
	// useful to include timestamps in logs. File names can be helpful though.
	l.SetFlags(logger.Lshortfile)
	ctx := logger.WithLogger(context.Background(), l)
	ctx, cancel := signal.NotifyContext(ctx, syscall.SIGTERM, syscall.SIGINT)
	defer cancel()
	if err := execute(ctx); err != nil {
		logger.Fatalf(ctx, err.Error())
	}
}

func execute(ctx context.Context) error {
	if buildDir == "" {
		return fmt.Errorf("must specify a Fuchsia build output directory")
	}

	// The package manifests generated by the build all use paths relative
	// to the build directory, so testsharder should change its working
	// directory to buildDir.
	wd, err := os.Getwd()
	if err != nil {
		return err
	}
	if err := os.Chdir(buildDir); err != nil {
		return err
	}
	defer os.Chdir(wd)

	targetDuration := time.Duration(targetDurationSecs) * time.Second
	if targetTestCount > 0 && targetDuration > 0 {
		return fmt.Errorf("max-shard-size and target-duration-secs cannot both be set")
	}

	perTestTimeout := time.Duration(perTestTimeoutSecs) * time.Second

	m, err := build.NewModules(buildDir)
	if err != nil {
		return err
	}

	if err = testsharder.ValidateTests(m.TestSpecs(), m.Platforms()); err != nil {
		return err
	}

	opts := &testsharder.ShardOptions{
		Mode: mode,
		Tags: tags,
	}
	shards := testsharder.MakeShards(m.TestSpecs(), opts)

	testDurations := testsharder.NewTestDurationsMap(m.TestDurations())

	var modifiers []testsharder.TestModifier
	if modifiersPath != "" {
		modifiers, err = testsharder.LoadTestModifiers(modifiersPath)
		if err != nil {
			return err
		}
	}

	if affectedTestsPath != "" {
		affectedModifiers, err := testsharder.AffectedModifiers(m.TestSpecs(), affectedTestsPath, affectedTestsMaxAttempts, affectedTestsMultiplyThreshold)
		if err != nil {
			return err
		}
		modifiers = append(modifiers, affectedModifiers...)
	}

	shards, err = testsharder.ShardAffected(shards, modifiers, affectedOnly)
	if err != nil {
		return err
	}

	shards, err = testsharder.MultiplyShards(ctx, shards, modifiers, testDurations, targetDuration, targetTestCount)
	if err != nil {
		return err
	}

	shards = testsharder.WithTargetDuration(shards, targetDuration, targetTestCount, maxShardsPerEnvironment, testDurations)

	if hermeticDeps {
		for _, s := range shards {
			testsharder.AddImageDeps(s, m.Images(), pave)
			if err := s.CreatePackageRepo(); err != nil {
				return err
			}
		}
	}

	if err := testsharder.ExtractDeps(shards, m.BuildDir()); err != nil {
		return err
	}

	if realmLabel != "" {
		testsharder.ApplyRealmLabel(shards, realmLabel)
	}

	if perTestTimeout > 0 {
		testsharder.ApplyTestTimeouts(shards, perTestTimeout)
	}

	f := os.Stdout
	if outputFile != "" {
		var err error
		f, err = os.Create(outputFile)
		if err != nil {
			return fmt.Errorf("unable to create %s: %v", outputFile, err)
		}
		defer f.Close()
	}

	encoder := json.NewEncoder(f)
	encoder.SetIndent("", "  ")
	if err := encoder.Encode(&shards); err != nil {
		return fmt.Errorf("failed to encode shards: %v", err)
	}
	return nil
}