github.com/nya3jp/tast@v0.0.0-20230601000426-85c8e4d83a9b/src/go.chromium.org/tast/core/internal/planner/downloader.go (about)

     1  // Copyright 2022 The ChromiumOS Authors
     2  // Use of this source code is governed by a BSD-style license that can be
     3  // found in the LICENSE file.
     4  
     5  package planner
     6  
     7  import (
     8  	"context"
     9  
    10  	"go.chromium.org/tast/core/errors"
    11  	"go.chromium.org/tast/core/internal/devserver"
    12  	"go.chromium.org/tast/core/internal/extdata"
    13  	"go.chromium.org/tast/core/internal/protocol"
    14  )
    15  
    16  // downloader encapsulates the logic to download external data files.
    17  type downloader struct {
    18  	m *extdata.Manager
    19  
    20  	pcfg           *Config
    21  	cl             devserver.Client
    22  	beforeDownload func(context.Context)
    23  }
    24  
    25  func newDownloader(ctx context.Context, pcfg *Config) (*downloader, error) {
    26  	cl, err := devserver.NewClient(
    27  		ctx,
    28  		pcfg.Service.GetDevservers(),
    29  		pcfg.Service.GetTlwServer(),
    30  		pcfg.Service.GetTlwSelfName(),
    31  		pcfg.Service.GetDutServer(),
    32  		pcfg.Service.GetSwarmingTaskID(),
    33  		pcfg.Service.GetBuildBucketID(),
    34  	)
    35  	if err != nil {
    36  		return nil, errors.Wrapf(err, "failed to create new client [devservers=%v, TLWServer=%s]",
    37  			pcfg.Service.GetDevservers(), pcfg.Service.GetTlwServer())
    38  	}
    39  	m, err := extdata.NewManager(ctx, pcfg.Dirs.GetDataDir(), pcfg.DataFile.GetBuildArtifactsUrl())
    40  	if err != nil {
    41  		return nil, err
    42  	}
    43  	return &downloader{
    44  		m:              m,
    45  		pcfg:           pcfg,
    46  		cl:             cl,
    47  		beforeDownload: pcfg.BeforeDownload,
    48  	}, nil
    49  }
    50  
    51  // TearDown must be called when downloader is destructed.
    52  func (d *downloader) TearDown() error {
    53  	return d.cl.TearDown()
    54  }
    55  
    56  // BeforeRun must be called before running a set of tests. It downloads external
    57  // data files if Config.DownloadMode is DownloadBatch.
    58  func (d *downloader) BeforeRun(ctx context.Context, tests []*protocol.Entity) {
    59  	if d.pcfg.DataFile.GetDownloadMode() == protocol.DownloadMode_BATCH {
    60  		// Ignore release because no data files are to be purged.
    61  		d.download(ctx, tests)
    62  	}
    63  }
    64  
    65  // BeforeEntity must be called before running each entity. It downloads external
    66  // data files if Config.DownloadMode is DownloadLazy.
    67  //
    68  // release must be called after entity finishes.
    69  func (d *downloader) BeforeEntity(ctx context.Context, entity *protocol.Entity) (release func()) {
    70  	if d.pcfg.DataFile.GetDownloadMode() == protocol.DownloadMode_LAZY {
    71  		return d.download(ctx, []*protocol.Entity{entity})
    72  	}
    73  	return func() {}
    74  }
    75  
    76  func (d *downloader) download(ctx context.Context, entities []*protocol.Entity) (release func()) {
    77  	jobs, release := d.m.PrepareDownloads(ctx, entities)
    78  	if len(jobs) > 0 {
    79  		if d.beforeDownload != nil {
    80  			d.beforeDownload(ctx)
    81  		}
    82  		extdata.RunDownloads(ctx, d.pcfg.Dirs.GetDataDir(), jobs, d.cl)
    83  	}
    84  	return release
    85  }