downloader: Allow downloads to actually be aborted

Previously the calls waiting for a download to be finish could be aborted but the downloads themselves would continue.
sdg0919 · Sep 26, 2016 · f376b9d · f376b9d
1 parent 4a88d70
commit f376b9d
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 4 deletions.
diff --git a/downloader/download.go b/downloader/download.go
@@ -125,13 +125,13 @@ func (dl *Download) download(req Request) (filename string, err error) {
 		}
 	}()
 
-	reader, err := dl.openBlob(req.URL)
+	blobReader, err := dl.openBlob(req.URL)
 	if err != nil {
 		return "", errors.Trace(err)
 	}
-	defer reader.Close()
+	defer blobReader.Close()
 
-	// XXX this should honor the Abort channel
+	reader := &abortableReader{blobReader, req.Abort}
 	_, err = io.Copy(tempFile, reader)
 	if err != nil {
 		return "", errors.Trace(err)
@@ -140,6 +140,23 @@ func (dl *Download) download(req Request) (filename string, err error) {
 	return tempFile.Name(), nil
 }
 
+// abortableReader wraps a Reader, returning an error from Read calls
+// if the abort channel provided is closed.
+type abortableReader struct {
+	r     io.Reader
+	abort <-chan struct{}
+}
+
+// Read implements io.Reader.
+func (ar *abortableReader) Read(p []byte) (int, error) {
+	select {
+	case <-ar.abort:
+		return 0, errors.New("download aborted")
+	default:
+	}
+	return ar.r.Read(p)
+}
+
 func verifyDownload(filename string, req Request) error {
 	if req.Verify == nil {
 		return nil

diff --git a/downloader/download_test.go b/downloader/download_test.go
@@ -139,6 +139,25 @@ func (s *DownloadSuite) TestVerifyInvalid(c *gc.C) {
 	checkDirEmpty(c, tmp)
 }
 
+func (s *DownloadSuite) TestAbort(c *gc.C) {
+	tmp := c.MkDir()
+	gitjujutesting.Server.Response(200, nil, []byte("archive"))
+	abort := make(chan struct{})
+	close(abort)
+	dl := downloader.StartDownload(
+		downloader.Request{
+			URL:       s.URL(c, "/archive.tgz"),
+			TargetDir: tmp,
+			Abort:     abort,
+		},
+		downloader.NewHTTPBlobOpener(utils.VerifySSLHostnames),
+	)
+	filename, err := dl.Wait()
+	c.Check(filename, gc.Equals, "")
+	c.Check(err, gc.ErrorMatches, ".+ download aborted")
+	checkDirEmpty(c, tmp)
+}
+
 func assertFileContents(c *gc.C, filename, expect string) {
 	got, err := ioutil.ReadFile(filename)
 	c.Assert(err, jc.ErrorIsNil)

diff --git a/worker/uniter/charm/bundles_test.go b/worker/uniter/charm/bundles_test.go
@@ -136,7 +136,7 @@ func (s *BundlesDirSuite) TestGet(c *gc.C) {
 
 	ch, err = d.Read(apiCharm, abort)
 	c.Assert(ch, gc.IsNil)
-	c.Assert(err, gc.ErrorMatches, regexp.QuoteMeta(`failed to download charm "cs:quantal/dummy-1" from API server: aborted`))
+	c.Assert(err, gc.ErrorMatches, regexp.QuoteMeta(`failed to download charm "cs:quantal/dummy-1" from API server: download aborted`))
 }
 
 func assertCharm(c *gc.C, bun charm.Bundle, sch *state.Charm) {