cmd/artifacts/copy.go

// Copyright 2019 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"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/google/subcommands"
	"go.chromium.org/luci/auth/client/authcli"
	"go.chromium.org/luci/hardcoded/chromeinfra"
	"go.fuchsia.dev/infra/artifacts"
	"go.fuchsia.dev/infra/buildbucket"
	"golang.org/x/sync/errgroup"
)

var stdout io.Writer = os.Stdout

type artifactsClient interface {
	GetBuildDir(bucket, build string) artifacts.Directory
}

type CopyCommand struct {
	authFlags authcli.Flags

	build string
	// The remote filepath with the target build's Cloud Storage directory.
	source string

	// The local path to write the artifact to.
	dest string

	// Whether source should be relative to the Cloud Storage bucket's root directory.
	// If false, source will be relative to the target build's Cloud Storage directory.
	fromRoot bool

	// The file containing the source paths to download.
	srcsFile string

	// The maximum number of concurrent downloading processes.
	j int

	// Whether to print verbose logs.
	verbose bool
}

func (*CopyCommand) Name() string {
	return "cp"
}

func (*CopyCommand) Usage() string {
	return "cp [flags...]"
}

func (*CopyCommand) Synopsis() string {
	return "fetches an artifact produced by a Fuchsia builder"
}

func (cmd *CopyCommand) SetFlags(f *flag.FlagSet) {
	cmd.authFlags.Register(flag.CommandLine, chromeinfra.DefaultAuthOptions())
	f.StringVar(&cmd.build, "build", "", "the ID of the build that produced the artifacts")
	f.StringVar(&cmd.source, "src", "", "The artifact file or directory to download from the build's Cloud Storage directory")
	f.StringVar(&cmd.dest, "dst", "", "The local path to write the artifact(s) to")
	f.BoolVar(&cmd.fromRoot, "root", false, "Whether src is relative to the root directory of the Cloud Storage bucket."+
		"If false, src will be taken as a relative path to the build-specific directory under the Cloud Storage bucket.")
	f.StringVar(&cmd.srcsFile, "srcs-file", "", "The file containing the source paths of the artifacts to download. These should be listed one path per line.")
	f.IntVar(&cmd.j, "j", 30, "The maximum number of concurrent downloading processes.")
	f.BoolVar(&cmd.verbose, "v", false, "Whether to print verbose logs.")
}

func (cmd *CopyCommand) Execute(ctx context.Context, f *flag.FlagSet, _ ...any) subcommands.ExitStatus {
	if err := cmd.validateAndExecute(ctx); err != nil {
		log.Println(err)
		return subcommands.ExitFailure
	}
	return subcommands.ExitSuccess
}

func (cmd *CopyCommand) validateAndExecute(ctx context.Context) error {
	opts, err := cmd.authFlags.Options()
	if err != nil {
		return err
	}

	if cmd.source == "" && cmd.srcsFile == "" {
		return fmt.Errorf("must provide at least one of -src or -srcs-file")
	}

	if cmd.dest == "" {
		return fmt.Errorf("missing -dst")
	}

	buildsCli, err := buildbucket.NewBuildsClient(ctx, buildbucket.DefaultHost, opts)
	if err != nil {
		return fmt.Errorf("failed to create builds client: %w", err)
	}

	artifactsCli, err := artifacts.NewClient(ctx, opts)
	if err != nil {
		return fmt.Errorf("failed to create artifacts client: %w", err)
	}

	return cmd.execute(ctx, buildsCli, artifactsCli)
}

type download struct {
	src  string
	dest string
}

func (cmd *CopyCommand) execute(ctx context.Context, buildsCli buildsClient, artifactsCli artifactsClient) error {
	bucket, err := getStorageBucket(ctx, buildsCli, cmd.build)
	if err != nil {
		return err
	}

	build := cmd.build
	if cmd.fromRoot {
		build = ""
	}
	dir := artifactsCli.GetBuildDir(bucket, build)

	var sourceList []download
	if cmd.srcsFile != "" {
		sourceFiles, err := os.ReadFile(cmd.srcsFile)
		if err != nil {
			return fmt.Errorf("failed to read src file %s: %w", cmd.srcsFile, err)
		}
		for _, src := range strings.Split(string(sourceFiles), "\n") {
			if src == "" {
				continue
			}
			sourceList = append(sourceList, cmd.constructDownloads(cmd.source, []string{src})...)
		}
	} else {
		objs, err := dir.List(ctx, cmd.source)
		if err != nil {
			return err
		}
		sourceList = append(sourceList, cmd.constructDownloads(cmd.source, objs)...)
	}

	eg, ctx := errgroup.WithContext(ctx)
	queueDownloads := make(chan download, len(sourceList))
	eg.Go(func() error {
		defer func() {
			close(queueDownloads)
		}()
		for _, s := range sourceList {
			queueDownloads <- s
		}
		return nil
	})

	var mu sync.Mutex
	var downloadedFileSizes []int64
	downloadFile := func() error {
		for download := range queueDownloads {
			startTime := time.Now()
			size, err := dir.CopyFile(ctx, download.src, download.dest)
			duration := time.Now().Sub(startTime)
			logStr := fmt.Sprintf("%s (%d bytes) to %s in %s", download.src, size, download.dest, duration)
			if err != nil {
				logStr = fmt.Sprintf("failed to download %s: %s", logStr, err)
			}
			mu.Lock()
			if cmd.verbose || err != nil {
				fmt.Fprintln(stdout, logStr)
			}
			downloadedFileSizes = append(downloadedFileSizes, size)
			mu.Unlock()
			if err != nil {
				return err
			}
		}
		return nil
	}
	startTime := time.Now()
	for i := 0; i < cmd.j; i++ {
		eg.Go(downloadFile)
	}
	err = eg.Wait()
	totalDuration := time.Now().Sub(startTime)
	var totalSize int64
	for _, size := range downloadedFileSizes {
		totalSize += size
	}
	fmt.Fprintf(stdout, "Num files: %d\nTotal bytes: %d\nDuration: %s\nSpeed: %.03fMB/s\n",
		len(downloadedFileSizes), totalSize, totalDuration, float64(totalSize)/(1000000*totalDuration.Seconds()))
	return err
}

func (cmd *CopyCommand) constructDownloads(src string, list []string) []download {
	var downloads []download
	for _, obj := range list {
		relPath := strings.TrimPrefix(obj, src)
		dest := filepath.Join(cmd.dest, relPath)
		downloads = append(downloads, download{src: obj, dest: dest})
	}
	return downloads
}