diff --git a/dfdaemon/downloader/dfget/dfget.go b/dfdaemon/downloader/dfget/dfget.go index 221d5eede..9165dcc95 100644 --- a/dfdaemon/downloader/dfget/dfget.go +++ b/dfdaemon/downloader/dfget/dfget.go @@ -17,6 +17,7 @@ package dfget import ( + "context" "fmt" netUrl "net/url" "os/exec" @@ -42,11 +43,11 @@ func NewGetter(cfg config.DFGetConfig) *DFGetter { return &DFGetter{config: cfg} } -// Download is the method of DFGetter to download by dragonfly. -func (dfGetter *DFGetter) Download(url string, header map[string][]string, name string) (string, error) { +// DownloadContext downloads the resources as specified in url. +func (dfGetter *DFGetter) DownloadContext(ctx context.Context, url string, header map[string][]string, name string) (string, error) { startTime := time.Now() dstPath := filepath.Join(dfGetter.config.DFRepo, name) - cmd := dfGetter.getCommand(url, header, dstPath) + cmd := dfGetter.getCommand(ctx, url, header, dstPath) err := cmd.Run() if cmd.ProcessState.Success() { log.Infof("dfget url:%s [SUCCESS] cost:%.3fs", url, time.Since(startTime).Seconds()) @@ -62,7 +63,7 @@ func (dfGetter *DFGetter) Download(url string, header map[string][]string, name // getCommand returns the command to download the given resource. func (dfGetter *DFGetter) getCommand( - url string, header map[string][]string, output string, + ctx context.Context, url string, header map[string][]string, output string, ) (cmd *exec.Cmd) { args := []string{ "-u", url, @@ -107,5 +108,5 @@ func (dfGetter *DFGetter) getCommand( } } - return exec.Command(dfGetter.config.DFPath, args...) + return exec.CommandContext(ctx, dfGetter.config.DFPath, args...) } diff --git a/dfdaemon/downloader/downloader.go b/dfdaemon/downloader/downloader.go index cada2978c..b5e4184dd 100644 --- a/dfdaemon/downloader/downloader.go +++ b/dfdaemon/downloader/downloader.go @@ -16,11 +16,13 @@ package downloader +import "context" + // Interface specifies on how an plugin can download a file. type Interface interface { - // Download download url file to file name - // return dst path and download error - Download(url string, header map[string][]string, name string) (string, error) + // DownloadContext downloads the resource as specified in url, and it accepts + // a context parameter so that it can handle timeouts correctly. + DownloadContext(ctx context.Context, url string, header map[string][]string, name string) (string, error) } // Factory is a function that returns a new downloader. diff --git a/dfdaemon/transport/transport.go b/dfdaemon/transport/transport.go index 349efa05b..d725c9743 100644 --- a/dfdaemon/transport/transport.go +++ b/dfdaemon/transport/transport.go @@ -17,6 +17,7 @@ package transport import ( + "context" "crypto/tls" "net" "net/http" @@ -133,7 +134,7 @@ func (roundTripper *DFRoundTripper) RoundTrip(req *http.Request) (*http.Response // download uses dfget to download. func (roundTripper *DFRoundTripper) download(req *http.Request, urlString string) (*http.Response, error) { - dstPath, err := roundTripper.downloadByGetter(urlString, req.Header, uuid.New()) + dstPath, err := roundTripper.downloadByGetter(req.Context(), urlString, req.Header, uuid.New()) if err != nil { logrus.Errorf("download fail: %v", err) return nil, err @@ -155,9 +156,9 @@ func (roundTripper *DFRoundTripper) download(req *http.Request, urlString string } // downloadByGetter is used to download file by DFGetter. -func (roundTripper *DFRoundTripper) downloadByGetter(url string, header map[string][]string, name string) (string, error) { +func (roundTripper *DFRoundTripper) downloadByGetter(ctx context.Context, url string, header map[string][]string, name string) (string, error) { logrus.Infof("start download url:%s to %s in repo", url, name) - return roundTripper.Downloader.Download(url, header, name) + return roundTripper.Downloader.DownloadContext(ctx, url, header, name) } // needUseGetter is the default value for ShouldUseDfget, which downloads all