implement download, upload

Author: adrienthebo

go.mod

@@ -35,6 +35,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/config v1.18.8 // indirect
 	github.com/aws/aws-sdk-go-v2/credentials v1.13.8 // indirect
 	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.21 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.11.47 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.27 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.21 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/ini v1.3.28 // indirect
@@ -55,6 +56,7 @@ require (
 	github.com/go-sourcemap/sourcemap v2.1.3+incompatible // indirect
 	github.com/godbus/dbus/v5 v5.1.0 // indirect
 	github.com/inconshreveable/mousetrap v1.0.1 // indirect
+	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/moby/sys/mountinfo v0.6.2 // indirect
 	github.com/opencontainers/selinux v1.10.2 // indirect
 	github.com/pkg/errors v0.9.1 // indirect

go.sum

@@ -9,6 +9,8 @@ github.com/aws/aws-sdk-go-v2/credentials v1.13.8 h1:vTrwTvv5qAwjWIGhZDSBH/oQHuIQ
 github.com/aws/aws-sdk-go-v2/credentials v1.13.8/go.mod h1:lVa4OHbvgjVot4gmh1uouF1ubgexSCN92P6CJQpT0t8=
 github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.21 h1:j9wi1kQ8b+e0FBVHxCqCGo4kxDU175hoDHcWAi0sauU=
 github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.21/go.mod h1:ugwW57Z5Z48bpvUyZuaPy4Kv+vEfJWnIrky7RmkBvJg=
+github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.11.47 h1:E884ndKWVGt8IhtUuGhXbEsmaCvdAAkTTUDu7uAok1g=
+github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.11.47/go.mod h1:KybsEsmXLO0u75FyS3F0sY4OQ97syDe8z+ISq8oEczA=
 github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.27 h1:I3cakv2Uy1vNmmhRQmFptYDxOvBnwCdNwyw63N0RaRU=
 github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.27/go.mod h1:a1/UpzeyBBerajpnP5nGZa9mGzsBn5cOKxm6NWQsvoI=
 github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.21 h1:5NbbMrIzmUn/TXFqAle6mgrH5m9cOvMLRGL7pnG8tRE=
@@ -88,6 +90,7 @@ github.com/in-toto/in-toto-golang v0.3.3 h1:tkkEBU5i09UEeWKnrp6Rq4fXKAfpVXYMLRO5
 github.com/in-toto/in-toto-golang v0.3.3/go.mod h1:dbXecHGZSqRubmm5TXtvDSZT5JyaKD7ebVTiC2aMLWY=
 github.com/inconshreveable/mousetrap v1.0.1 h1:U3uMjPSQEBMNp1lFxmllqCPM6P5u/Xq7Pgzkat/bFNc=
 github.com/inconshreveable/mousetrap v1.0.1/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
+github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
 github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
 github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
 github.com/karrick/godirwalk v1.17.0 h1:b4kY7nqDdioR/6qnbHQyDvmA17u5G1cZ6J+CZXwSWoI=

pkg/leeway/cache.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
 	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/aws/aws-sdk-go-v2/service/s3/types"
 	"github.com/aws/aws-sdk-go-v2/service/sts"
@@ -250,10 +251,10 @@ type S3RemoteCache struct {
 func NewS3RemoteCache(bucketName string, cfg *aws.Config) (*S3RemoteCache, error) {
 	if cfg == nil {
 		v, err := config.LoadDefaultConfig(context.TODO())
-		cfg = &v
 		if err != nil {
 			return nil, fmt.Errorf("cannot load s3 config: %s", err)
 		}
+		cfg = &v
 	}
 	s3Client := s3.NewFromConfig(*cfg)
 
@@ -287,7 +288,7 @@ func (rs *S3RemoteCache) ExistingPackages(pkgs []*Package) (map[*Package]struct{
 			continue
 		}
 
-		packagesToKeys[p] = fmt.Sprintf("%s.tar.gz", version)
+		packagesToKeys[p] = filepath.Base(fmt.Sprintf("%s.tar.gz", version))
 	}
 
 	if len(packagesToKeys) == 0 {
@@ -296,12 +297,12 @@ func (rs *S3RemoteCache) ExistingPackages(pkgs []*Package) (map[*Package]struct{
 	log.Debugf("Checking if %d packages exist in the remote cache using s3", len(packagesToKeys))
 
 	ch := make(chan *Package, len(packagesToKeys))
-	defer close(ch)
 
 	existingPackages := make(map[*Package]struct{})
 	wg := sync.WaitGroup{}
 
 	for pkg, key := range packagesToKeys {
+		wg.Add(1)
 		go func(pkg *Package, key string) {
 			defer wg.Done()
 
@@ -311,25 +312,108 @@ func (rs *S3RemoteCache) ExistingPackages(pkgs []*Package) (map[*Package]struct{
 				ch <- pkg
 			}
 		}(pkg, key)
-
-		wg.Add(1)
 	}
 	wg.Wait()
+	close(ch)
+
+	for p := range ch {
+		existingPackages[p] = struct{}{}
+	}
 
 	return existingPackages, nil
 }
 
 // Download makes a best-effort attempt at downloading previously cached build artifacts for the given packages
 // in their current version. A cache miss (i.e. a build artifact not being available) does not constitute an
 // error. Get should try and download as many artifacts as possible.
-func (s3 *S3RemoteCache) Download(dst Cache, pkgs []*Package) error {
-	panic("not implemented") // TODO: Implement
+func (rs *S3RemoteCache) Download(dst Cache, pkgs []*Package) error {
+	fmt.Printf("☁️  downloading %d cached build artifacts from s3 remote cache\n", len(pkgs))
+	var (
+		files []string
+		dest  string
+	)
+
+	for _, pkg := range pkgs {
+		fn, exists := dst.Location(pkg)
+		if exists {
+			continue
+		}
+
+		if dest == "" {
+			dest = filepath.Dir(fn)
+		} else if dest != filepath.Dir(fn) {
+			return xerrors.Errorf("s3 cache only supports one target folder, not %s and %s", dest, filepath.Dir(fn))
+		}
+
+		files = append(files, fmt.Sprintf("%s/%s", rs.BucketName, strings.TrimLeft(fn, "/")))
+	}
+
+	wg := sync.WaitGroup{}
+
+	for _, file := range files {
+		wg.Add(1)
+
+		go func(file string) {
+			defer wg.Done()
+
+			key := filepath.Base(file)
+			fields := log.Fields{
+				"key":    key,
+				"bucket": rs.BucketName,
+				"region": rs.s3Config.Region,
+			}
+			log.WithFields(fields).Debug("downloading object from s3")
+			len, err := rs.getObject(context.TODO(), key, fmt.Sprintf("%s/%s", dest, key))
+			if err != nil {
+				log.WithFields(fields).Warnf("failed to download and store object %s from s3: %s", key, err)
+			} else {
+				log.WithFields(fields).Debugf("downloaded %d byte object from s3 to %s", len, file)
+			}
+
+		}(file)
+	}
+	wg.Wait()
+
+	return nil
 }
 
 // Upload makes a best effort to upload the build arfitacts to a remote cache. If uploading an artifact fails, that
 // does not constitute an error.
-func (s3 *S3RemoteCache) Upload(src Cache, pkgs []*Package) error {
-	panic("not implemented") // TODO: Implement
+func (rs *S3RemoteCache) Upload(src Cache, pkgs []*Package) error {
+	var files []string
+	for _, pkg := range pkgs {
+		file, exists := src.Location(pkg)
+		if !exists {
+			continue
+		}
+		files = append(files, file)
+	}
+	fmt.Fprintf(os.Stdout, "☁️  uploading %d build artifacts to s3 remote cache\n", len(files))
+
+	wg := sync.WaitGroup{}
+
+	for _, file := range files {
+		wg.Add(1)
+		go func(file string) {
+			defer wg.Done()
+			key := filepath.Base(file)
+			fields := log.Fields{
+				"key":    key,
+				"bucket": rs.BucketName,
+				"region": rs.s3Config.Region,
+			}
+			log.WithFields(fields).Debug("uploading object to s3")
+			res, err := rs.uploadObject(context.TODO(), filepath.Base(file), file)
+			if err != nil {
+				log.WithFields(fields).Warnf("Failed to upload object to s3: %s", err)
+			} else {
+				log.WithFields(fields).Debugf("completed upload to %s", res.Location)
+			}
+		}(file)
+	}
+	wg.Wait()
+
+	return nil
 }
 
 func (rs *S3RemoteCache) hasBucket(ctx context.Context) (bool, error) {
@@ -356,11 +440,10 @@ func (rs *S3RemoteCache) hasBucket(ctx context.Context) (bool, error) {
 }
 
 func (rs *S3RemoteCache) hasObject(ctx context.Context, key string) (bool, error) {
-	cfg := *rs.s3Config
 	fields := log.Fields{
 		"key":    key,
 		"bucket": rs.BucketName,
-		"region": cfg.Region,
+		"region": rs.s3Config.Region,
 	}
 	log.WithFields(fields).Debugf("Checking s3 for cached package")
 
@@ -377,7 +460,10 @@ func (rs *S3RemoteCache) hasObject(ctx context.Context, key string) (bool, error
 		}
 
 		// We've received an error that's not a simple missing key error. Collect more information
-		_, _ = rs.hasBucket(ctx)
+		hasBucket, _ := rs.hasBucket(ctx)
+		if !hasBucket {
+			return false, err
+		}
 
 		log.WithFields(fields).Warnf("S3 GetObject failed: %s", err)
 		return false, err
@@ -386,3 +472,47 @@ func (rs *S3RemoteCache) hasObject(ctx context.Context, key string) (bool, error
 	// XXX
 	return true, nil
 }
+
+func (rs *S3RemoteCache) getObject(ctx context.Context, key string, path string) (int64, error) {
+
+	var partMiBs int64 = 10
+	downloader := manager.NewDownloader(rs.s3Client, func(d *manager.Downloader) {
+		d.PartSize = partMiBs * 1024 * 1024
+	})
+	buffer := manager.NewWriteAtBuffer([]byte{})
+	res, err := downloader.Download(context.TODO(), buffer, &s3.GetObjectInput{
+		Bucket: aws.String(rs.BucketName),
+		Key:    aws.String(key),
+	})
+	if err != nil {
+		return res, err
+	}
+
+	err = os.WriteFile(path, buffer.Bytes(), 0644)
+	if err != nil {
+		return 0, xerrors.Errorf("failed to write s3 download to %s: %s", path, err)
+	}
+	return res, nil
+}
+
+func (rs *S3RemoteCache) uploadObject(ctx context.Context, key string, path string) (*manager.UploadOutput, error) {
+
+	fN, err := os.Open(path)
+	if err != nil {
+		return nil, xerrors.Errorf("cannot open %s for S3 upload: %s", path, err)
+	}
+
+	var partMiBs int64 = 10
+	uploader := manager.NewUploader(rs.s3Client, func(u *manager.Uploader) {
+		u.PartSize = partMiBs * 1024 * 1024
+	})
+	res, err := uploader.Upload(context.TODO(), &s3.PutObjectInput{
+		Bucket: aws.String(rs.BucketName),
+		Key:    aws.String(key),
+		Body:   fN,
+	})
+	if err != nil {
+		return res, err
+	}
+	return res, nil
+}