s3: upload package files concurrently

context/context.go

@@ -436,11 +436,11 @@ func (context *AptlyContext) GetPublishedStorage(name string) (aptly.PublishedSt
 			}
 
 			var err error
-			publishedStorage, err = s3.NewPublishedStorage(
+			publishedStorage, err = s3.NewPublishedStorageWithUploadConcurrency(
 				params.AccessKeyID, params.SecretAccessKey, params.SessionToken,
 				params.Region, params.Endpoint, params.Bucket, params.ACL, params.Prefix, params.StorageClass,
 				params.EncryptionMethod, params.PlusWorkaround, params.DisableMultiDel,
-				params.ForceSigV2, params.ForceVirtualHostedStyle, params.Debug)
+				params.ForceSigV2, params.ForceVirtualHostedStyle, params.UploadConcurrency, params.Debug)
 			if err != nil {
 				return nil, err
 			}

deb/publish.go

@@ -829,6 +829,32 @@ func (p *PublishedRepo) GetSkelFiles(skelDir string, component string) (map[stri
 	return files, nil
 }
 
+func (p *PublishedRepo) packagePublishPath(pkg *Package, component string) (string, bool, error) {
+	for _, arch := range p.Architectures {
+		if !pkg.MatchesArchitecture(arch) {
+			continue
+		}
+
+		if pkg.IsInstaller {
+			if p.Distribution == aptly.DistributionFocal {
+				return filepath.Join("dists", p.Distribution, component, fmt.Sprintf("%s-%s", pkg.Name, arch), "current", "legacy-images"), true, nil
+			}
+			return filepath.Join("dists", p.Distribution, component, fmt.Sprintf("%s-%s", pkg.Name, arch), "current", "images"), true, nil
+		}
+
+		poolDir, err := pkg.PoolDirectory()
+		if err != nil {
+			return "", false, err
+		}
+		if p.MultiDist {
+			return filepath.Join("pool", p.Distribution, component, poolDir), true, nil
+		}
+		return filepath.Join("pool", component, poolDir), true, nil
+	}
+
+	return "", false, nil
+}
+
 // Publish publishes snapshot (repository) contents, links package files, generates Packages & Release files, signs them
 func (p *PublishedRepo) Publish(packagePool aptly.PackagePool, publishedStorageProvider aptly.PublishedStorageProvider,
 	collectionFactory *CollectionFactory, signer pgp.Signer, progress aptly.Progress, forceOverwrite bool, skelDir string) error {
@@ -919,6 +945,13 @@ func (p *PublishedRepo) Publish(packagePool aptly.PackagePool, publishedStorageP
 		progress.InitBar(count, false, aptly.BarPublishGeneratePackageFiles)
 	}
 
+	uploadConcurrency := 1
+	if uploader, ok := publishedStorage.(interface{ UploadConcurrency() int }); ok {
+		if concurrency := uploader.UploadConcurrency(); concurrency > 1 {
+			uploadConcurrency = min(concurrency, 128)
+		}
+	}
+
 	for component, list := range lists {
 		hadUdebs := false
 
@@ -928,43 +961,24 @@ func (p *PublishedRepo) Publish(packagePool aptly.PackagePool, publishedStorageP
 		}
 
 		list.PrepareIndex()
-
 		contentIndexes := map[string]*ContentsIndex{}
 
-		err = list.ForEachIndexed(func(pkg *Package) error {
+		processPackage := func(pkg *Package) error {
 			if progress != nil {
 				progress.AddBar(1)
 			}
 
-			for _, arch := range p.Architectures {
-				if pkg.MatchesArchitecture(arch) {
-					hadUdebs = hadUdebs || pkg.IsUdeb
-
-					var relPath string
-					if !pkg.IsInstaller {
-						poolDir, err2 := pkg.PoolDirectory()
-						if err2 != nil {
-							return err2
-						}
-						if p.MultiDist {
-							relPath = filepath.Join("pool", p.Distribution, component, poolDir)
-						} else {
-							relPath = filepath.Join("pool", component, poolDir)
-						}
-
-					} else {
-						if p.Distribution == aptly.DistributionFocal {
-							relPath = filepath.Join("dists", p.Distribution, component, fmt.Sprintf("%s-%s", pkg.Name, arch), "current", "legacy-images")
-						} else {
-							relPath = filepath.Join("dists", p.Distribution, component, fmt.Sprintf("%s-%s", pkg.Name, arch), "current", "images")
-						}
-					}
-
+			relPath, matches, err2 := p.packagePublishPath(pkg, component)
+			if err2 != nil {
+				return err2
+			}
+			if matches {
+				hadUdebs = hadUdebs || pkg.IsUdeb
+				if uploadConcurrency == 1 {
 					err = pkg.LinkFromPool(publishedStorage, packagePool, p.Prefix, relPath, forceOverwrite)
 					if err != nil {
 						return err
 					}
-					break
 				}
 			}
 
@@ -1017,7 +1031,70 @@ func (p *PublishedRepo) Publish(packagePool aptly.PackagePool, publishedStorageP
 			pkg.contents = nil
 
 			return batch.Write()
-		})
+		}
+
+		if uploadConcurrency == 1 {
+			err = list.ForEachIndexed(processPackage)
+		} else {
+			for offset := 0; offset < len(list.packagesIndex); offset += uploadConcurrency {
+				batch := list.packagesIndex[offset:min(offset+uploadConcurrency, len(list.packagesIndex))]
+				groups := map[string][]*Package{}
+				var paths []string
+				for _, pkg := range batch {
+					relPath, matches, err2 := p.packagePublishPath(pkg, component)
+					if err2 != nil {
+						return fmt.Errorf("unable to prepare package uploads: %s", err2)
+					}
+					if !matches {
+						continue
+					}
+
+					files := pkg.Files()
+					for i := range files {
+						if _, err2 = files[i].GetPoolPath(packagePool); err2 != nil {
+							return fmt.Errorf("unable to prepare package uploads: %s", err2)
+						}
+					}
+					if pkg.IsSource {
+						pkg.Extra()
+					}
+					if _, exists := groups[relPath]; !exists {
+						paths = append(paths, relPath)
+					}
+					groups[relPath] = append(groups[relPath], pkg)
+				}
+
+				done := make(chan struct{}, len(paths))
+				uploadErrors := make([]error, len(paths))
+				for i, relPath := range paths {
+					go func() {
+						defer func() { done <- struct{}{} }()
+						for _, pkg := range groups[relPath] {
+							if err := pkg.LinkFromPool(publishedStorage, packagePool, p.Prefix, relPath, forceOverwrite); err != nil {
+								uploadErrors[i] = err
+								return
+							}
+						}
+					}()
+				}
+				for range paths {
+					<-done
+				}
+				for _, uploadErr := range uploadErrors {
+					if uploadErr != nil {
+						return fmt.Errorf("unable to upload packages: %s", uploadErr)
+					}
+				}
+				for _, pkg := range batch {
+					if err = processPackage(pkg); err != nil {
+						break
+					}
+				}
+				if err != nil {
+					break
+				}
+			}
+		}
 
 		if err != nil {
 			return fmt.Errorf("unable to process packages: %s", err)

deb/publish_test.go

@@ -8,6 +8,8 @@ import (
 	"os"
 	"path/filepath"
 	"sort"
+	"sync"
+	"time"
 
 	"github.com/aptly-dev/aptly/aptly"
 	"github.com/aptly-dev/aptly/database"
@@ -62,6 +64,49 @@ type FakeStorageProvider struct {
 	storages map[string]aptly.PublishedStorage
 }
 
+type concurrentTestStorage struct {
+	aptly.PublishedStorage
+	concurrency      int
+	mu               sync.Mutex
+	active           int
+	maxActive        int
+	release          chan struct{}
+	releaseTriggered bool
+}
+
+func (s *concurrentTestStorage) UploadConcurrency() int {
+	return s.concurrency
+}
+
+func (s *concurrentTestStorage) LinkFromPool(publishedPrefix, publishedRelPath, fileName string,
+	sourcePool aptly.PackagePool, sourcePath string, sourceChecksums utils.ChecksumInfo, force bool) error {
+	if s.release == nil {
+		return s.PublishedStorage.LinkFromPool(publishedPrefix, publishedRelPath, fileName, sourcePool, sourcePath, sourceChecksums, force)
+	}
+
+	s.mu.Lock()
+	s.active++
+	s.maxActive = max(s.maxActive, s.active)
+	if s.active == s.concurrency && !s.releaseTriggered {
+		close(s.release)
+		s.releaseTriggered = true
+	}
+	s.mu.Unlock()
+	defer func() {
+		s.mu.Lock()
+		s.active--
+		s.mu.Unlock()
+	}()
+
+	select {
+	case <-s.release:
+	case <-time.After(time.Second):
+		return errors.New("package uploads did not overlap")
+	}
+
+	return s.PublishedStorage.LinkFromPool(publishedPrefix, publishedRelPath, fileName, sourcePool, sourcePath, sourceChecksums, force)
+}
+
 func (p *FakeStorageProvider) GetPublishedStorage(name string) (aptly.PublishedStorage, error) {
 	storage, ok := p.storages[name]
 	if !ok {
@@ -195,6 +240,8 @@ func (s *PublishedRepoSuite) TestNewPublishedRepo(c *C) {
 }
 
 func (s *PublishedRepoSuite) TestMultiDistPool(c *C) {
+	s.provider.storages[""] = &concurrentTestStorage{PublishedStorage: s.publishedStorage, concurrency: 3}
+
 	repo, err := NewPublishedRepo("", "ppa", "squeeze", nil, []string{"main"}, []interface{}{s.snapshot}, s.factory, true)
 	c.Assert(err, IsNil)
 	err = repo.Publish(s.packagePool, s.provider, s.factory, &NullSigner{}, nil, false, "")
@@ -246,6 +293,24 @@ func (s *PublishedRepoSuite) TestMultiDistPool(c *C) {
 
 }
 
+func (s *PublishedRepoSuite) TestPackageUploadsConcurrent(c *C) {
+	for _, pkg := range []*Package{s.p1, s.p2, s.p3} {
+		pkg.Source = ""
+		c.Assert(s.packageCollection.Update(pkg), IsNil)
+	}
+
+	storage := &concurrentTestStorage{
+		PublishedStorage: s.publishedStorage,
+		concurrency:      2,
+		release:          make(chan struct{}),
+	}
+	s.provider.storages[""] = storage
+
+	err := s.repo.Publish(s.packagePool, s.provider, s.factory, &NullSigner{}, nil, false, "")
+	c.Assert(err, IsNil)
+	c.Check(storage.maxActive, Equals, 2)
+}
+
 func (s *PublishedRepoSuite) TestPrefixNormalization(c *C) {
 
 	for _, t := range []struct {

debian/aptly.conf

@@ -281,6 +281,10 @@ s3_publish_endpoints:
     #     # Disable path style visit, useful with non-AWS S3-compatible object stores
     #     # which only support virtual hosted style
     #     force_virtualhosted_style: false
+    #     # Upload Concurrency (optional)
+    #     # Maximum number of package files uploaded concurrently (up to 128), defaults to 4.
+    #     # Set to 1 to upload package files sequentially.
+    #     upload_concurrency: 4
     #     # Debug (optional)
     #     # Enables detailed request/response dump for each S3 operation
     #     debug: false
@@ -424,4 +428,3 @@ packagepool_storage:
     # # See: Azure documentation https://docs.microsoft.com/en-us/azure/storage/common/storage-configure-connection-string
     # # defaults to "https://<accountName>.blob.core.windows.net"
     # endpoint: ""
-