| // 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 |
| } |