Allow for setting concurrency to a non-default value

[#163240667] `kiln fetch` can download releases in parallel

Signed-off-by: Slawek Ligus <sligus@pivotal.io>

Author: ciriarte

acceptance/help_test.go

@@ -61,6 +61,7 @@ Usage: kiln [options] fetch [<args>]
 
 Command Arguments:
   --assets-file, -a          string (required)  path to assets file
+  --download-threads, -dt    int                number of parallel threads to download parts from S3
   --releases-directory, -rd  string (required)  path to a directory to download releases into
 `
 

commands/fetch.go

@@ -40,10 +40,11 @@ type Fetch struct {
 	S3Provider func(*session.Session, ...*aws.Config) s3iface.S3API
 
 	Options struct {
-		AssetsFile     string   `short:"a" long:"assets-file" required:"true" description:"path to assets file"`
-		VariablesFiles []string `short:"vf" long:"variables-file" description:"path to variables file"`
-		Variables      []string `short:"vr" long:"variable" description:"variable in key=value format"`
-		ReleasesDir    string   `short:"rd" long:"releases-directory" required:"true" description:"path to a directory to download releases into"`
+		AssetsFile      string   `short:"a" long:"assets-file" required:"true" description:"path to assets file"`
+		VariablesFiles  []string `short:"vf" long:"variables-file" description:"path to variables file"`
+		Variables       []string `short:"vr" long:"variable" description:"variable in key=value format"`
+		ReleasesDir     string   `short:"rd" long:"releases-directory" required:"true" description:"path to a directory to download releases into"`
+		DownloadThreads int      `short:"dt" long:"download-threads" description:"number of parallel threads to download parts from S3"`
 	}
 }
 
@@ -135,10 +136,16 @@ type Downloader interface {
 	Download(w io.WriterAt, input *s3.GetObjectInput, options ...func(*s3manager.Downloader)) (n int64, err error)
 }
 
-func DownloadReleases(logger *log.Logger, assetsLock cargo.AssetsLock, bucket string, matchedS3Objects map[cargo.CompiledRelease]string, fileCreator func(string) (io.WriterAt, error), downloader Downloader) error {
+func DownloadReleases(logger *log.Logger, assetsLock cargo.AssetsLock, bucket string, matchedS3Objects map[cargo.CompiledRelease]string, fileCreator func(string) (io.WriterAt, error), downloader Downloader, downloadThreads int) error {
 	releases := assetsLock.Releases
 	stemcell := assetsLock.Stemcell
 
+	setConcurrency := func(dl *s3manager.Downloader) {
+		if downloadThreads > 0 {
+			dl.Concurrency = downloadThreads
+		}
+	}
+
 	for _, release := range releases {
 		s3Key, ok := matchedS3Objects[cargo.CompiledRelease{
 			Name:            release.Name,
@@ -161,7 +168,7 @@ func DownloadReleases(logger *log.Logger, assetsLock cargo.AssetsLock, bucket st
 		_, err = downloader.Download(file, &s3.GetObjectInput{
 			Bucket: aws.String(bucket),
 			Key:    aws.String(s3Key),
-		})
+		}, setConcurrency)
 
 		if err != nil {
 			return fmt.Errorf("failed to download file, %v\n", err)
@@ -241,7 +248,7 @@ func (f Fetch) Execute(args []string) error {
 	}
 
 	downloader := s3manager.NewDownloaderWithClient(s3Client)
-	return DownloadReleases(f.logger, assetsLock, assets.CompiledReleases.Bucket, MatchedS3Objects, fileCreator, downloader)
+	return DownloadReleases(f.logger, assetsLock, assets.CompiledReleases.Bucket, MatchedS3Objects, fileCreator, downloader, f.Options.DownloadThreads)
 }
 
 func (f Fetch) Usage() jhanda.Usage {

commands/fetch_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/s3"
 	"github.com/aws/aws-sdk-go/service/s3/s3iface"
+	"github.com/aws/aws-sdk-go/service/s3/s3manager"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 	"github.com/pivotal-cf/jhanda"
@@ -34,10 +35,7 @@ compiled_releases:
   regex: ^2.5/.+/(?P<release_name>[a-z-_]+)-(?P<release_version>[0-9\.]+)-(?P<stemcell_os>[a-z-_]+)-(?P<stemcell_version>[\d\.]+)\.tgz$
 `
 
-const MinimalAssetsLockContents = `
----
-releases: []
-`
+const MinimalAssetsLockContents = `---`
 
 var _ = Describe("Fetch", func() {
 	var (
@@ -163,6 +161,7 @@ compiled_releases:
 				Expect(sessionArg.Config.Region).To(Equal(aws.String("north-east-1")))
 			})
 		})
+
 		Context("failure cases", func() {
 			Context("the assets-file flag is missing", func() {
 				It("returns a flag error", func() {
@@ -320,25 +319,56 @@ compiled_releases:
 		})
 
 		It("downloads the appropriate versions of releases listed in the assets.lock", func() {
-			err = commands.DownloadReleases(logger, assetsLock, bucket, matchedS3Objects, fileCreator, fakeDownloader)
+			err = commands.DownloadReleases(logger, assetsLock, bucket, matchedS3Objects, fileCreator, fakeDownloader, 7)
 			Expect(err).NotTo(HaveOccurred())
 			Expect(fakeDownloader.DownloadCallCount()).To(Equal(2))
 
-			w1, input1, _ := fakeDownloader.DownloadArgsForCall(0)
+			w1, input1, opts := fakeDownloader.DownloadArgsForCall(0)
 			Expect(w1).To(Equal(fakeUAAFile))
 			Expect(input1).To(Equal(uaaInput))
+			Expect(opts).To(HaveLen(1))
+
+			downloader := &s3manager.Downloader{
+				Concurrency: s3manager.DefaultDownloadConcurrency,
+			}
+
+			opts[0](downloader)
 
-			w2, input2, _ := fakeDownloader.DownloadArgsForCall(1)
+			Expect(downloader.Concurrency).To(Equal(7))
+
+			w2, input2, opts := fakeDownloader.DownloadArgsForCall(1)
 			Expect(w2).To(Equal(fakeBPMFile))
 			Expect(input2).To(Equal(bpmInput))
+			Expect(opts).To(HaveLen(1))
+		})
+
+		Context("when number of threads is not specified", func() {
+			It("uses the s3manager package's default download concurrency", func() {
+				err = commands.DownloadReleases(logger, assetsLock, bucket, matchedS3Objects, fileCreator, fakeDownloader, 0)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(fakeDownloader.DownloadCallCount()).To(Equal(2))
+
+				w1, input1, opts := fakeDownloader.DownloadArgsForCall(0)
+				Expect(w1).To(Equal(fakeUAAFile))
+				Expect(input1).To(Equal(uaaInput))
+				Expect(opts).To(HaveLen(1))
+
+				downloader := &s3manager.Downloader{
+					Concurrency: s3manager.DefaultDownloadConcurrency,
+				}
+
+				opts[0](downloader)
+
+				Expect(downloader.Concurrency).To(Equal(s3manager.DefaultDownloadConcurrency))
+			})
 		})
 
 		It("returns an error if the release does not exist", func() {
 			assetsLock.Releases = []cargo.Release{
 				{Name: "not-real", Version: "1.2.3"},
 			}
 
-			err = commands.DownloadReleases(logger, assetsLock, bucket, matchedS3Objects, fileCreator, fakeDownloader)
+			err = commands.DownloadReleases(logger, assetsLock, bucket, matchedS3Objects, fileCreator, fakeDownloader, 0)
 			Expect(err).To(HaveOccurred())
 			Expect(err).To(MatchError("Compiled release: not-real, version: 1.2.3, stemcell OS: ubuntu-trusty, stemcell version: 1234, not found"))
 		})