gitaly-bench.go

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"github.com/pkg/profile"
	"gitlab.com/gitlab-org/gitaly-bench/benchmarks"
	"gitlab.com/gitlab-org/gitaly-bench/operations"
	"gitlab.com/gitlab-org/gitaly-proto/go/gitalypb"
	gitalyauth "gitlab.com/gitlab-org/gitaly/auth"
	"gitlab.com/gitlab-org/gitaly/client"
	"google.golang.org/grpc"
)

var (
	hostFlag     = flag.String("host", "tcp://127.0.0.1:9999", "The server address in the format of host:port")
	storage      = flag.String("storage", "default", "The storage to test")
	gitRepoPath  = flag.String("repo", "", "Relative path of git repo within storage")
	iterations   = flag.Int("iterations", 10, "Number of iterations to perform on each thread")
	parallel     = flag.Int("concurrency", 10, "Number of concurrent threads to test the benchmark on")
	profileFlag  = flag.String("profile", "", "Profiling mode: `block` / `cpu` / `mem` / `mutex` ")
	authFlag     = flag.String("auth", "", "Authentication token")
	authFileFlag = flag.String("auth-file", "", "Path to a text file containing the authentication token")
)

func newConnection() (*grpc.ClientConn, error) {
	connOpts := make([]grpc.DialOption, len(client.DefaultDialOpts))
	copy(connOpts, client.DefaultDialOpts)

	token := ""

	if *authFileFlag != "" {
		contents, err := ioutil.ReadFile(*authFileFlag)
		if err != nil {
			log.Fatalf("unable to read auth-file: %v", err)
		}
		token = string(contents)
		token = strings.TrimRight(token, "\n\r")
	} else if *authFlag != "" {
		token = *authFlag
	}

	if token != "" {
		connOpts = append(connOpts, grpc.WithPerRPCCredentials(gitalyauth.RPCCredentialsV2(token)))
	}

	conn, err := client.Dial(*hostFlag, connOpts)
	if err != nil {
		return nil, err
	}

	return conn, nil
}

func getProfileOpt(name string) func(p *profile.Profile) {
	switch *profileFlag {
	case "block":
		return profile.BlockProfile
	case "cpu":
		return profile.CPUProfile
	case "mem":
		return profile.MemProfile
	case "mutex":
		return profile.MutexProfile
	default:
		return nil
	}
}

func main() {
	commands := map[string]*struct {
		flagSet *flag.FlagSet
		builder func(repository *gitalypb.Repository, conn *grpc.ClientConn) operations.BenchmarkOperation
	}{
		"gc":                    {flagSet: operations.GCFlagSet, builder: operations.NewGC},
		"repack-full":           {flagSet: operations.GCFlagSet, builder: operations.NewRepackFull},
		"find-all-branch-names": {flagSet: nil, builder: operations.NewFindAllBranchNames},
		"find-all-branches":     {flagSet: nil, builder: operations.NewFindAllBranches},
		"commit-diff":           {flagSet: operations.CommitDiffFlagSet, builder: operations.NewCommitDiff},
		"ref-exists":            {flagSet: operations.RefExistsFlagSet, builder: operations.NewRefExists},
		"info-refs":             {flagSet: nil, builder: operations.NewInfoRefs},
		"has-local-branches":    {flagSet: nil, builder: operations.NewHasLocalBranches},
		"find-commit":           {flagSet: operations.FindCommitFlagSet, builder: operations.NewFindCommit},
	}

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of gitaly-bench [options] command:\n\n")
		fmt.Fprintf(os.Stderr, "Options:\n")
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n`command` being one of:\n")
		for k := range commands {
			fmt.Fprintf(os.Stderr, "\t%v\n", k)
		}
		fmt.Fprintf(os.Stderr, "Use gitaly-bench command --help for command options\n")
	}

	flag.Parse()

	if *gitRepoPath == "" {
		flag.Usage()
		os.Exit(1)
	}

	// Verify that a subcommand has been provided
	if len(flag.Args()) < 1 {
		flag.Usage()
		os.Exit(1)
	}

	if *profileFlag != "" {
		profileOpt := getProfileOpt(*profileFlag)
		if profileOpt == nil {
			flag.Usage()
			os.Exit(1)
		}

		defer profile.Start(profileOpt, profile.ProfilePath(".")).Stop()
	}

	benchmarkName := flag.Arg(0)
	command := commands[benchmarkName]
	if command == nil {
		flag.Usage()
		os.Exit(1)
	}

	if command.flagSet != nil {
		err := command.flagSet.Parse(flag.Args()[1:])
		if err != nil || len(command.flagSet.Args()) > 0 {
			if len(command.flagSet.Args()) > 0 {
				fmt.Println("Unknown argument", strings.Join(command.flagSet.Args(), " "))
			}
			command.flagSet.Usage()
			os.Exit(1)
		}
	} else {
		if len(flag.Args()) > 1 {
			fmt.Println("Unknown argument", strings.Join(flag.Args()[1:], " "))
			flag.Usage()
			os.Exit(1)
		}
	}

	conn, err := newConnection()
	if err != nil {
		panic(err)
	}

	defer conn.Close()

	repository := &gitalypb.Repository{
		StorageName:  *storage,
		RelativePath: *gitRepoPath,
	}

	operation := command.builder(repository, conn)
	benchmarks.RunBenchmark(benchmarkName, *parallel, *iterations, operation)
}