Use Go context to abort gracefully mirror updates

There are two fixes here: 1. Abort package download immediately as ^C is pressed. 2. Import all the already downloaded files into package pool, so that next time mirror is updated, aptly won't download them once again.
2026-01-11 03:11:50 +00:00 · 2017-11-30 00:49:37 +03:00
parent a037615962
commit 15618c8ea8
12 changed files with 131 additions and 80 deletions
--- a/http/compression.go
+++ b/http/compression.go
@@ -3,6 +3,7 @@ package http
 import (
 	"compress/bzip2"
 	"compress/gzip"
+	"context"
 	"fmt"
 	"io"
 	"net/url"
@@ -39,7 +40,7 @@ var compressionMethods = []struct {

 // DownloadTryCompression tries to download from URL .bz2, .gz and raw extension until
 // it finds existing file.
-func DownloadTryCompression(downloader aptly.Downloader, baseURL *url.URL, path string, expectedChecksums map[string]utils.ChecksumInfo, ignoreMismatch bool, maxTries int) (io.Reader, *os.File, error) {
+func DownloadTryCompression(ctx context.Context, downloader aptly.Downloader, baseURL *url.URL, path string, expectedChecksums map[string]utils.ChecksumInfo, ignoreMismatch bool, maxTries int) (io.Reader, *os.File, error) {
 	var err error

 	for _, method := range compressionMethods {
@@ -63,13 +64,13 @@ func DownloadTryCompression(downloader aptly.Downloader, baseURL *url.URL, path

 		if foundChecksum {
 			expected := expectedChecksums[bestSuffix]
-			file, err = DownloadTempWithChecksum(downloader, tryURL.String(), &expected, ignoreMismatch, maxTries)
+			file, err = DownloadTempWithChecksum(ctx, downloader, tryURL.String(), &expected, ignoreMismatch, maxTries)
 		} else {
 			if !ignoreMismatch {
 				continue
 			}

-			file, err = DownloadTemp(downloader, tryURL.String())
+			file, err = DownloadTemp(ctx, downloader, tryURL.String())
 		}

 		if err != nil {
--- a/http/compression_test.go
+++ b/http/compression_test.go
@@ -1,6 +1,7 @@
 package http

 import (
+	"context"
 	"errors"
 	"io"
 	"net/url"
@@ -12,6 +13,7 @@ import (

 type CompressionSuite struct {
 	baseURL *url.URL
+	ctx     context.Context
 }

 var _ = Suite(&CompressionSuite{})
@@ -25,6 +27,7 @@ const (

 func (s *CompressionSuite) SetUpTest(c *C) {
 	s.baseURL, _ = url.Parse("http://example.com/")
+	s.ctx = context.Background()
 }

 func (s *CompressionSuite) TestDownloadTryCompression(c *C) {
@@ -41,7 +44,7 @@ func (s *CompressionSuite) TestDownloadTryCompression(c *C) {
 	buf = make([]byte, 4)
 	d := NewFakeDownloader()
 	d.ExpectResponse("http://example.com/file.bz2", bzipData)
-	r, file, err := DownloadTryCompression(d, s.baseURL, "file", expectedChecksums, false, 1)
+	r, file, err := DownloadTryCompression(s.ctx, d, s.baseURL, "file", expectedChecksums, false, 1)
 	c.Assert(err, IsNil)
 	defer file.Close()
 	io.ReadFull(r, buf)
@@ -53,7 +56,7 @@ func (s *CompressionSuite) TestDownloadTryCompression(c *C) {
 	d = NewFakeDownloader()
 	d.ExpectError("http://example.com/file.bz2", &Error{Code: 404})
 	d.ExpectResponse("http://example.com/file.gz", gzipData)
-	r, file, err = DownloadTryCompression(d, s.baseURL, "file", expectedChecksums, false, 1)
+	r, file, err = DownloadTryCompression(s.ctx, d, s.baseURL, "file", expectedChecksums, false, 1)
 	c.Assert(err, IsNil)
 	defer file.Close()
 	io.ReadFull(r, buf)
@@ -66,7 +69,7 @@ func (s *CompressionSuite) TestDownloadTryCompression(c *C) {
 	d.ExpectError("http://example.com/file.bz2", &Error{Code: 404})
 	d.ExpectError("http://example.com/file.gz", &Error{Code: 404})
 	d.ExpectResponse("http://example.com/file.xz", xzData)
-	r, file, err = DownloadTryCompression(d, s.baseURL, "file", expectedChecksums, false, 1)
+	r, file, err = DownloadTryCompression(s.ctx, d, s.baseURL, "file", expectedChecksums, false, 1)
 	c.Assert(err, IsNil)
 	defer file.Close()
 	io.ReadFull(r, buf)
@@ -80,7 +83,7 @@ func (s *CompressionSuite) TestDownloadTryCompression(c *C) {
 	d.ExpectError("http://example.com/file.gz", &Error{Code: 404})
 	d.ExpectError("http://example.com/file.xz", &Error{Code: 404})
 	d.ExpectResponse("http://example.com/file", rawData)
-	r, file, err = DownloadTryCompression(d, s.baseURL, "file", expectedChecksums, false, 1)
+	r, file, err = DownloadTryCompression(s.ctx, d, s.baseURL, "file", expectedChecksums, false, 1)
 	c.Assert(err, IsNil)
 	defer file.Close()
 	io.ReadFull(r, buf)
@@ -91,7 +94,7 @@ func (s *CompressionSuite) TestDownloadTryCompression(c *C) {
 	d = NewFakeDownloader()
 	d.ExpectError("http://example.com/file.bz2", &Error{Code: 404})
 	d.ExpectResponse("http://example.com/file.gz", "x")
-	_, _, err = DownloadTryCompression(d, s.baseURL, "file", nil, true, 1)
+	_, _, err = DownloadTryCompression(s.ctx, d, s.baseURL, "file", nil, true, 1)
 	c.Assert(err, ErrorMatches, "unexpected EOF")
 	c.Assert(d.Empty(), Equals, true)
 }
@@ -109,7 +112,7 @@ func (s *CompressionSuite) TestDownloadTryCompressionLongestSuffix(c *C) {
 	buf = make([]byte, 4)
 	d := NewFakeDownloader()
 	d.ExpectResponse("http://example.com/subdir/file.bz2", bzipData)
-	r, file, err := DownloadTryCompression(d, s.baseURL, "subdir/file", expectedChecksums, false, 1)
+	r, file, err := DownloadTryCompression(s.ctx, d, s.baseURL, "subdir/file", expectedChecksums, false, 1)
 	c.Assert(err, IsNil)
 	defer file.Close()
 	io.ReadFull(r, buf)
@@ -119,7 +122,7 @@ func (s *CompressionSuite) TestDownloadTryCompressionLongestSuffix(c *C) {

 func (s *CompressionSuite) TestDownloadTryCompressionErrors(c *C) {
 	d := NewFakeDownloader()
-	_, _, err := DownloadTryCompression(d, s.baseURL, "file", nil, true, 1)
+	_, _, err := DownloadTryCompression(s.ctx, d, s.baseURL, "file", nil, true, 1)
 	c.Assert(err, ErrorMatches, "unexpected request.*")

 	d = NewFakeDownloader()
@@ -127,7 +130,7 @@ func (s *CompressionSuite) TestDownloadTryCompressionErrors(c *C) {
 	d.ExpectError("http://example.com/file.gz", &Error{Code: 404})
 	d.ExpectError("http://example.com/file.xz", &Error{Code: 404})
 	d.ExpectError("http://example.com/file", errors.New("403"))
-	_, _, err = DownloadTryCompression(d, s.baseURL, "file", nil, true, 1)
+	_, _, err = DownloadTryCompression(s.ctx, d, s.baseURL, "file", nil, true, 1)
 	c.Assert(err, ErrorMatches, "403")

 	d = NewFakeDownloader()
@@ -141,6 +144,6 @@ func (s *CompressionSuite) TestDownloadTryCompressionErrors(c *C) {
 		"file.xz":  {Size: 7},
 		"file":     {Size: 7},
 	}
-	_, _, err = DownloadTryCompression(d, s.baseURL, "file", expectedChecksums, false, 1)
+	_, _, err = DownloadTryCompression(s.ctx, d, s.baseURL, "file", expectedChecksums, false, 1)
 	c.Assert(err, ErrorMatches, "checksums don't match.*")
 }
--- a/http/download.go
+++ b/http/download.go
@@ -1,12 +1,15 @@
 package http

 import (
+	"context"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"os"
 	"path/filepath"
 	"strings"
+	"syscall"
 	"time"

 	"github.com/mxk/go-flowrate/flowrate"
@@ -62,12 +65,24 @@ func (downloader *downloaderImpl) GetProgress() aptly.Progress {
 }

 // Download starts new download task
-func (downloader *downloaderImpl) Download(url string, destination string) error {
-	return downloader.DownloadWithChecksum(url, destination, nil, false, 1)
+func (downloader *downloaderImpl) Download(ctx context.Context, url string, destination string) error {
+	return downloader.DownloadWithChecksum(ctx, url, destination, nil, false, 1)
+}
+
+func retryableError(err error) bool {
+	switch err.(type) {
+	case net.Error:
+		return true
+	case *net.OpError:
+		return true
+	case syscall.Errno:
+		return true
+	}
+	return false
 }

 // DownloadWithChecksum starts new download task with checksum verification
-func (downloader *downloaderImpl) DownloadWithChecksum(url string, destination string,
+func (downloader *downloaderImpl) DownloadWithChecksum(ctx context.Context, url string, destination string,
 	expected *utils.ChecksumInfo, ignoreMismatch bool, maxTries int) error {

 	downloader.progress.Printf("Downloading %s...\n", url)
@@ -77,6 +92,7 @@ func (downloader *downloaderImpl) DownloadWithChecksum(url string, destination s
 		return errors.Wrap(err, url)
 	}
 	req.Close = true
+	req = req.WithContext(ctx)

 	proxyURL, _ := downloader.client.Transport.(*http.Transport).Proxy(req)
 	if proxyURL == nil && (req.URL.Scheme == "http" || req.URL.Scheme == "https") {
@@ -88,10 +104,10 @@ func (downloader *downloaderImpl) DownloadWithChecksum(url string, destination s
 	for maxTries > 0 {
 		temppath, err = downloader.download(req, url, destination, expected, ignoreMismatch)

-		if err != nil {
+		if err != nil && retryableError(err) {
 			maxTries--
 		} else {
-			// successful download
+			// get out of the loop
 			break
 		}
 	}
--- a/http/download_test.go
+++ b/http/download_test.go
@@ -1,6 +1,7 @@
 package http

 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"net"
@@ -21,6 +22,7 @@ type DownloaderSuiteBase struct {
 	ch       chan struct{}
 	progress aptly.Progress
 	d        aptly.Downloader
+	ctx      context.Context
 }

 func (s *DownloaderSuiteBase) SetUpTest(c *C) {
@@ -44,6 +46,7 @@ func (s *DownloaderSuiteBase) SetUpTest(c *C) {
 	s.progress.Start()

 	s.d = NewDownloader(0, s.progress)
+	s.ctx = context.Background()
 }

 func (s *DownloaderSuiteBase) TearDownTest(c *C) {
@@ -71,52 +74,52 @@ func (s *DownloaderSuite) TearDownTest(c *C) {
 }

 func (s *DownloaderSuite) TestDownloadOK(c *C) {
-	c.Assert(s.d.Download(s.url+"/test", s.tempfile.Name()), IsNil)
+	c.Assert(s.d.Download(s.ctx, s.url+"/test", s.tempfile.Name()), IsNil)
 }

 func (s *DownloaderSuite) TestDownloadWithChecksum(c *C) {
-	c.Assert(s.d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{}, false, 1),
+	c.Assert(s.d.DownloadWithChecksum(s.ctx, s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{}, false, 1),
 		ErrorMatches, ".*size check mismatch 12 != 0")

-	c.Assert(s.d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{Size: 12, MD5: "abcdef"}, false, 1),
+	c.Assert(s.d.DownloadWithChecksum(s.ctx, s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{Size: 12, MD5: "abcdef"}, false, 1),
 		ErrorMatches, ".*md5 hash mismatch \"a1acb0fe91c7db45ec4d775192ec5738\" != \"abcdef\"")

-	c.Assert(s.d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{Size: 12, MD5: "abcdef"}, true, 1),
+	c.Assert(s.d.DownloadWithChecksum(s.ctx, s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{Size: 12, MD5: "abcdef"}, true, 1),
 		IsNil)

-	c.Assert(s.d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738"}, false, 1),
+	c.Assert(s.d.DownloadWithChecksum(s.ctx, s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738"}, false, 1),
 		IsNil)

-	c.Assert(s.d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738", SHA1: "abcdef"}, false, 1),
+	c.Assert(s.d.DownloadWithChecksum(s.ctx, s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738", SHA1: "abcdef"}, false, 1),
 		ErrorMatches, ".*sha1 hash mismatch \"921893bae6ad6fd818401875d6779254ef0ff0ec\" != \"abcdef\"")

-	c.Assert(s.d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738",
+	c.Assert(s.d.DownloadWithChecksum(s.ctx, s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738",
 		SHA1: "921893bae6ad6fd818401875d6779254ef0ff0ec"}, false, 1),
 		IsNil)

-	c.Assert(s.d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738",
+	c.Assert(s.d.DownloadWithChecksum(s.ctx, s.url+"/test", s.tempfile.Name(), &utils.ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738",
 		SHA1: "921893bae6ad6fd818401875d6779254ef0ff0ec", SHA256: "abcdef"}, false, 1),
 		ErrorMatches, ".*sha256 hash mismatch \"b3c92ee1246176ed35f6e8463cd49074f29442f5bbffc3f8591cde1dcc849dac\" != \"abcdef\"")

 	checksums := utils.ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738",
 		SHA1: "921893bae6ad6fd818401875d6779254ef0ff0ec", SHA256: "b3c92ee1246176ed35f6e8463cd49074f29442f5bbffc3f8591cde1dcc849dac"}
-	c.Assert(s.d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), &checksums, false, 1),
+	c.Assert(s.d.DownloadWithChecksum(s.ctx, s.url+"/test", s.tempfile.Name(), &checksums, false, 1),
 		IsNil)
 	// download backfills missing checksums
 	c.Check(checksums.SHA512, Equals, "bac18bf4e564856369acc2ed57300fecba3a2c1af5ae8304021e4252488678feb18118466382ee4e1210fe1f065080210e453a80cfb37ccb8752af3269df160e")
 }

 func (s *DownloaderSuite) TestDownload404(c *C) {
-	c.Assert(s.d.Download(s.url+"/doesntexist", s.tempfile.Name()),
+	c.Assert(s.d.Download(s.ctx, s.url+"/doesntexist", s.tempfile.Name()),
 		ErrorMatches, "HTTP code 404.*")
 }

 func (s *DownloaderSuite) TestDownloadConnectError(c *C) {
-	c.Assert(s.d.Download("http://nosuch.localhost/", s.tempfile.Name()),
+	c.Assert(s.d.Download(s.ctx, "http://nosuch.localhost/", s.tempfile.Name()),
 		ErrorMatches, ".*no such host")
 }

 func (s *DownloaderSuite) TestDownloadFileError(c *C) {
-	c.Assert(s.d.Download(s.url+"/test", "/"),
+	c.Assert(s.d.Download(s.ctx, s.url+"/test", "/"),
 		ErrorMatches, ".*permission denied")
 }
--- a/http/fake.go
+++ b/http/fake.go
@@ -1,6 +1,7 @@
 package http

 import (
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -60,7 +61,7 @@ func (f *FakeDownloader) Empty() bool {
 }

 // DownloadWithChecksum performs fake download by matching against first expectation in the queue or any expectation, with cheksum verification
-func (f *FakeDownloader) DownloadWithChecksum(url string, filename string, expected *utils.ChecksumInfo, ignoreMismatch bool, maxTries int) error {
+func (f *FakeDownloader) DownloadWithChecksum(ctx context.Context, url string, filename string, expected *utils.ChecksumInfo, ignoreMismatch bool, maxTries int) error {
 	var expectation expectedRequest
 	if len(f.expected) > 0 && f.expected[0].URL == url {
 		expectation, f.expected = f.expected[0], f.expected[1:]
@@ -109,8 +110,8 @@ func (f *FakeDownloader) DownloadWithChecksum(url string, filename string, expec
 }

 // Download performs fake download by matching against first expectation in the queue
-func (f *FakeDownloader) Download(url string, filename string) error {
-	return f.DownloadWithChecksum(url, filename, nil, false, 1)
+func (f *FakeDownloader) Download(ctx context.Context, url string, filename string) error {
+	return f.DownloadWithChecksum(ctx, url, filename, nil, false, 1)
 }

 // GetProgress returns Progress object
--- a/http/temp.go
+++ b/http/temp.go
@@ -1,6 +1,7 @@
 package http

 import (
+	"context"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -12,14 +13,14 @@ import (
 // DownloadTemp starts new download to temporary file and returns File
 //
 // Temporary file would be already removed, so no need to cleanup
-func DownloadTemp(downloader aptly.Downloader, url string) (*os.File, error) {
-	return DownloadTempWithChecksum(downloader, url, nil, false, 1)
+func DownloadTemp(ctx context.Context, downloader aptly.Downloader, url string) (*os.File, error) {
+	return DownloadTempWithChecksum(ctx, downloader, url, nil, false, 1)
 }

 // DownloadTempWithChecksum is a DownloadTemp with checksum verification
 //
 // Temporary file would be already removed, so no need to cleanup
-func DownloadTempWithChecksum(downloader aptly.Downloader, url string, expected *utils.ChecksumInfo, ignoreMismatch bool, maxTries int) (*os.File, error) {
+func DownloadTempWithChecksum(ctx context.Context, downloader aptly.Downloader, url string, expected *utils.ChecksumInfo, ignoreMismatch bool, maxTries int) (*os.File, error) {
 	tempdir, err := ioutil.TempDir(os.TempDir(), "aptly")
 	if err != nil {
 		return nil, err
@@ -33,7 +34,7 @@ func DownloadTempWithChecksum(downloader aptly.Downloader, url string, expected
 		defer downloader.GetProgress().ShutdownBar()
 	}

-	err = downloader.DownloadWithChecksum(url, tempfile, expected, ignoreMismatch, maxTries)
+	err = downloader.DownloadWithChecksum(ctx, url, tempfile, expected, ignoreMismatch, maxTries)
 	if err != nil {
 		return nil, err
 	}
--- a/http/temp_test.go
+++ b/http/temp_test.go
@@ -23,7 +23,7 @@ func (s *TempSuite) TearDownTest(c *C) {
 }

 func (s *TempSuite) TestDownloadTemp(c *C) {
-	f, err := DownloadTemp(s.d, s.url+"/test")
+	f, err := DownloadTemp(s.ctx, s.d, s.url+"/test")
 	c.Assert(err, IsNil)
 	defer f.Close()

@@ -37,18 +37,18 @@ func (s *TempSuite) TestDownloadTemp(c *C) {
 }

 func (s *TempSuite) TestDownloadTempWithChecksum(c *C) {
-	f, err := DownloadTempWithChecksum(s.d, s.url+"/test", &utils.ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738",
+	f, err := DownloadTempWithChecksum(s.ctx, s.d, s.url+"/test", &utils.ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738",
 		SHA1: "921893bae6ad6fd818401875d6779254ef0ff0ec", SHA256: "b3c92ee1246176ed35f6e8463cd49074f29442f5bbffc3f8591cde1dcc849dac"}, false, 1)
 	c.Assert(err, IsNil)

 	c.Assert(f.Close(), IsNil)

-	_, err = DownloadTempWithChecksum(s.d, s.url+"/test", &utils.ChecksumInfo{Size: 13}, false, 1)
+	_, err = DownloadTempWithChecksum(s.ctx, s.d, s.url+"/test", &utils.ChecksumInfo{Size: 13}, false, 1)
 	c.Assert(err, ErrorMatches, ".*size check mismatch 12 != 13")
 }

 func (s *TempSuite) TestDownloadTempError(c *C) {
-	f, err := DownloadTemp(s.d, s.url+"/doesntexist")
+	f, err := DownloadTemp(s.ctx, s.d, s.url+"/doesntexist")
 	c.Assert(err, NotNil)
 	c.Assert(f, IsNil)
 	c.Assert(err, ErrorMatches, "HTTP code 404.*")