feat(layers): add layer balancer script (#1643)

Co-authored-by: Leandro Damascena <[email protected]>

layer/scripts/layer-balancer/README.md

@@ -0,0 +1,37 @@
+<!-- markdownlint-disable MD041 MD043 -->
+# Layer balancer
+
+This folder contains a Go project that balances the layer version of Lambda Powertools across all regions, so
+every region has the same layer version.
+
+Before:
+
+```text
+arn:aws:lambda:eu-central-1:017000801446:layer:AWSLambdaPowertoolsPythonV2:11
+...
+arn:aws:lambda:eu-west-1:017000801446:layer:AWSLambdaPowertoolsPythonV2:9
+```
+
+After:
+
+```text
+arn:aws:lambda:eu-central-1:017000801446:layer:AWSLambdaPowertoolsPythonV2:11
+...
+arn:aws:lambda:eu-west-1:017000801446:layer:AWSLambdaPowertoolsPythonV2:11
+```
+
+## What's happening under the hood?
+
+1. Query all regions to find the greatest version number
+2. Download the latest layer from eu-central-1
+3. Use the layer contents to bump the version on each region until it matches 1
+
+## Requirements
+
+* go >= 1.18
+
+## How to use
+
+1. Set your AWS_PROFILE to the correct profile
+2. `go run .`
+3. Profit :-)

layer/scripts/layer-balancer/go.mod

@@ -0,0 +1,24 @@
+module layerbalancer
+
+go 1.18
+
+require (
+	github.com/aws/aws-sdk-go-v2 v1.16.16
+	github.com/aws/aws-sdk-go-v2/config v1.17.8
+	github.com/aws/aws-sdk-go-v2/service/lambda v1.24.6
+	golang.org/x/sync v0.1.0
+)
+
+require (
+	github.com/aws/aws-sdk-go-v2/credentials v1.12.21 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.17 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.23 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.17 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/ini v1.3.24 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.17 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sso v1.11.23 // indirect
+	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.13.6 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sts v1.16.19 // indirect
+	github.com/aws/smithy-go v1.13.3 // indirect
+	github.com/jmespath/go-jmespath v0.4.0 // indirect
+)

layer/scripts/layer-balancer/go.sum

@@ -0,0 +1,37 @@
+github.com/aws/aws-sdk-go-v2 v1.16.16 h1:M1fj4FE2lB4NzRb9Y0xdWsn2P0+2UHVxwKyOa4YJNjk=
+github.com/aws/aws-sdk-go-v2 v1.16.16/go.mod h1:SwiyXi/1zTUZ6KIAmLK5V5ll8SiURNUYOqTerZPaF9k=
+github.com/aws/aws-sdk-go-v2/config v1.17.8 h1:b9LGqNnOdg9vR4Q43tBTVWk4J6F+W774MSchvKJsqnE=
+github.com/aws/aws-sdk-go-v2/config v1.17.8/go.mod h1:UkCI3kb0sCdvtjiXYiU4Zx5h07BOpgBTtkPu/49r+kA=
+github.com/aws/aws-sdk-go-v2/credentials v1.12.21 h1:4tjlyCD0hRGNQivh5dN8hbP30qQhMLBE/FgQR1vHHWM=
+github.com/aws/aws-sdk-go-v2/credentials v1.12.21/go.mod h1:O+4XyAt4e+oBAoIwNUYkRg3CVMscaIJdmZBOcPgJ8D8=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.17 h1:r08j4sbZu/RVi+BNxkBJwPMUYY3P8mgSDuKkZ/ZN1lE=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.17/go.mod h1:yIkQcCDYNsZfXpd5UX2Cy+sWA1jPgIhGTw9cOBzfVnQ=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.23 h1:s4g/wnzMf+qepSNgTvaQQHNxyMLKSawNhKCPNy++2xY=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.23/go.mod h1:2DFxAQ9pfIRy0imBCJv+vZ2X6RKxves6fbnEuSry6b4=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.17 h1:/K482T5A3623WJgWT8w1yRAFK4RzGzEl7y39yhtn9eA=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.17/go.mod h1:pRwaTYCJemADaqCbUAxltMoHKata7hmB5PjEXeu0kfg=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.3.24 h1:wj5Rwc05hvUSvKuOF29IYb9QrCLjU+rHAy/x/o0DK2c=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.3.24/go.mod h1:jULHjqqjDlbyTa7pfM7WICATnOv+iOhjletM3N0Xbu8=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.17 h1:Jrd/oMh0PKQc6+BowB+pLEwLIgaQF29eYbe7E1Av9Ug=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.17/go.mod h1:4nYOrY41Lrbk2170/BGkcJKBhws9Pfn8MG3aGqjjeFI=
+github.com/aws/aws-sdk-go-v2/service/lambda v1.24.6 h1:N7RkXX2SJbN+TCp295J3LdMR0KRFd2Bhi5nIO+svLQY=
+github.com/aws/aws-sdk-go-v2/service/lambda v1.24.6/go.mod h1:oTJIIluTaJCRT6xP1AZpuU3JwRHBC0Q5O4Hg+SUxFHw=
+github.com/aws/aws-sdk-go-v2/service/sso v1.11.23 h1:pwvCchFUEnlceKIgPUouBJwK81aCkQ8UDMORfeFtW10=
+github.com/aws/aws-sdk-go-v2/service/sso v1.11.23/go.mod h1:/w0eg9IhFGjGyyncHIQrXtU8wvNsTJOP0R6PPj0wf80=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.13.6 h1:OwhhKc1P9ElfWbMKPIbMMZBV6hzJlL2JKD76wNNVzgQ=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.13.6/go.mod h1:csZuQY65DAdFBt1oIjO5hhBR49kQqop4+lcuCjf2arA=
+github.com/aws/aws-sdk-go-v2/service/sts v1.16.19 h1:9pPi0PsFNAGILFfPCk8Y0iyEBGc6lu6OQ97U7hmdesg=
+github.com/aws/aws-sdk-go-v2/service/sts v1.16.19/go.mod h1:h4J3oPZQbxLhzGnk+j9dfYHi5qIOVJ5kczZd658/ydM=
+github.com/aws/smithy-go v1.13.3 h1:l7LYxGuzK6/K+NzJ2mC+VvLUbae0sL3bXU//04MkmnA=
+github.com/aws/smithy-go v1.13.3/go.mod h1:Tg+OJXh4MB2R/uN61Ko2f6hTZwB/ZYGOtib8J3gBHzA=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+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/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
+golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

layer/scripts/layer-balancer/main.go

@@ -0,0 +1,292 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"io/ioutil"
+	"log"
+	"net/http"
+	"os"
+	"os/signal"
+	"sort"
+	"sync"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/lambda"
+	"github.com/aws/aws-sdk-go-v2/service/lambda/types"
+	"golang.org/x/sync/errgroup"
+)
+
+type LayerInfo struct {
+	Name         string
+	Description  string
+	Architecture types.Architecture
+
+	LayerContentOnce sync.Once
+	LayerContent     []byte
+}
+
+// canonicalLayers are the layers that we want to keep in sync across all regions
+var canonicalLayers = []LayerInfo{
+	{
+		Name:         "AWSLambdaPowertoolsPythonV2",
+		Description:  "Lambda Powertools for Python [x86_64] with extra dependencies version bump",
+		Architecture: types.ArchitectureX8664,
+	},
+	{
+		Name:         "AWSLambdaPowertoolsPythonV2-Arm64",
+		Description:  "Lambda Powertools for Python [arm64] with extra dependencies version bump",
+		Architecture: types.ArchitectureArm64,
+	},
+}
+
+// regions are the regions that we want to keep in sync
+var regions = []string{
+	"af-south-1",
+	"eu-central-1",
+	"us-east-1",
+	"us-east-2",
+	"us-west-1",
+	"us-west-2",
+	"ap-east-1",
+	"ap-south-1",
+	"ap-northeast-1",
+	"ap-northeast-2",
+	"ap-southeast-1",
+	"ap-southeast-2",
+	"ca-central-1",
+	"eu-west-1",
+	"eu-west-2",
+	"eu-west-3",
+	"eu-south-1",
+	"eu-north-1",
+	"sa-east-1",
+	"ap-southeast-3",
+	"ap-northeast-3",
+	"me-south-1",
+}
+
+// getLayerVersion returns the latest version of a layer in a region
+func getLayerVersion(ctx context.Context, layerName string, region string) (int64, error) {
+	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
+	if err != nil {
+		return 0, err
+	}
+
+	lambdaSvc := lambda.NewFromConfig(cfg)
+
+	layerVersionsResult, err := lambdaSvc.ListLayerVersions(ctx, &lambda.ListLayerVersionsInput{
+		LayerName: aws.String(layerName),
+		MaxItems:  aws.Int32(1),
+	})
+	if err != nil {
+		return 0, err
+	}
+
+	if len(layerVersionsResult.LayerVersions) == 0 {
+		return 0, fmt.Errorf("no layer meets the search criteria %s - %s", layerName, region)
+	}
+	return layerVersionsResult.LayerVersions[0].Version, nil
+}
+
+// getGreatestVersion returns the greatest version of a layer across all regions
+func getGreatestVersion(ctx context.Context) (int64, error) {
+	var versions []int64
+
+	g, ctx := errgroup.WithContext(ctx)
+
+	for idx := range canonicalLayers {
+		layer := &canonicalLayers[idx]
+
+		for _, region := range regions {
+			layerName := layer.Name
+			ctx := ctx
+			region := region
+
+			g.Go(func() error {
+				version, err := getLayerVersion(ctx, layerName, region)
+				if err != nil {
+					return err
+				}
+
+				log.Printf("[%s] %s -> %d", layerName, region, version)
+
+				versions = append(versions, version)
+				return nil
+			})
+		}
+	}
+
+	if err := g.Wait(); err != nil {
+		return 0, err
+	}
+
+	// Find the maximum version by reverse sorting the versions array
+	sort.Slice(versions, func(i, j int) bool { return versions[i] > versions[j] })
+	return versions[0], nil
+}
+
+// balanceRegionToVersion creates a new layer version in a region with the same contents as the canonical layer, until it matches the maxVersion
+func balanceRegionToVersion(ctx context.Context, region string, layer *LayerInfo, maxVersion int64) error {
+	currentLayerVersion, err := getLayerVersion(ctx, layer.Name, region)
+	if err != nil {
+		return fmt.Errorf("error getting layer version: %w", err)
+	}
+
+	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
+	if err != nil {
+		return err
+	}
+
+	lambdaSvc := lambda.NewFromConfig(cfg)
+
+	for i := currentLayerVersion; i < maxVersion; i++ {
+		log.Printf("[%s] Bumping %s to version %d (max %d)", layer.Name, region, i, maxVersion)
+
+		payload, err := downloadCanonicalLayerZip(ctx, layer)
+		if err != nil {
+			return fmt.Errorf("error downloading canonical zip: %w", err)
+		}
+
+		layerVersionResponse, err := lambdaSvc.PublishLayerVersion(ctx, &lambda.PublishLayerVersionInput{
+			Content: &types.LayerVersionContentInput{
+				ZipFile: payload,
+			},
+			LayerName:               aws.String(layer.Name),
+			CompatibleArchitectures: []types.Architecture{layer.Architecture},
+			CompatibleRuntimes:      []types.Runtime{types.RuntimePython37, types.RuntimePython38, types.RuntimePython39},
+			Description:             aws.String(layer.Description),
+			LicenseInfo:             aws.String("MIT-0"),
+		})
+		if err != nil {
+			return fmt.Errorf("error publishing layer version: %w", err)
+		}
+
+		_, err = lambdaSvc.AddLayerVersionPermission(ctx, &lambda.AddLayerVersionPermissionInput{
+			Action:        aws.String("lambda:GetLayerVersion"),
+			LayerName:     aws.String(layer.Name),
+			Principal:     aws.String("*"),
+			StatementId:   aws.String("PublicLayerAccess"),
+			VersionNumber: layerVersionResponse.Version,
+		})
+		if err != nil {
+			return fmt.Errorf("error making layer public: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// balanceRegions creates new layer versions in all regions with the same contents as the canonical layer, until they match the maxVersion
+func balanceRegions(ctx context.Context, maxVersion int64) error {
+	g, ctx := errgroup.WithContext(ctx)
+
+	for idx := range canonicalLayers {
+		layer := &canonicalLayers[idx]
+
+		for _, region := range regions {
+			ctx := ctx
+			region := region
+			layer := layer
+			version := maxVersion
+
+			g.Go(func() error {
+				return balanceRegionToVersion(ctx, region, layer, version)
+			})
+		}
+	}
+
+	if err := g.Wait(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// downloadCanonicalLayerZip downloads the canonical layer zip file that will be used to bump the versions later
+func downloadCanonicalLayerZip(ctx context.Context, layer *LayerInfo) ([]byte, error) {
+	var innerErr error
+
+	layer.LayerContentOnce.Do(func() {
+		// We use eu-central-1 as the canonical region to download the Layer from
+		cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion("eu-central-1"))
+		if err != nil {
+			innerErr = err
+		}
+
+		lambdaSvc := lambda.NewFromConfig(cfg)
+
+		// Gets the latest version of the layer
+		version, err := getLayerVersion(ctx, layer.Name, "eu-central-1")
+		if err != nil {
+			innerErr = fmt.Errorf("error getting eu-central-1 layer version: %w", err)
+		}
+
+		// Gets the Layer content URL from S3
+		getLayerVersionResult, err := lambdaSvc.GetLayerVersion(ctx, &lambda.GetLayerVersionInput{
+			LayerName:     aws.String(layer.Name),
+			VersionNumber: version,
+		})
+		if err != nil {
+			innerErr = fmt.Errorf("error getting eu-central-1 layer download URL: %w", err)
+		}
+
+		s3LayerUrl := getLayerVersionResult.Content.Location
+		log.Printf("[%s] Downloading Layer from %s", layer.Name, *s3LayerUrl)
+
+		resp, err := http.Get(*s3LayerUrl)
+		if err != nil {
+			innerErr = err
+		}
+		defer resp.Body.Close()
+
+		body, err := ioutil.ReadAll(resp.Body)
+		if err != nil {
+			innerErr = err
+		}
+
+		layer.LayerContent = body
+	})
+
+	return layer.LayerContent, innerErr
+}
+
+func main() {
+	ctx := context.Background()
+
+	// Cancel everything if interrupted
+	ctx, cancel := context.WithCancel(ctx)
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, os.Interrupt)
+	defer func() {
+		signal.Stop(c)
+		cancel()
+	}()
+	go func() {
+		select {
+		case <-c:
+			cancel()
+		case <-ctx.Done():
+		}
+	}()
+
+	// Find the greatest layer version across all regions
+	greatestVersion, err := getGreatestVersion(ctx)
+	if err != nil {
+		cancel()
+		log.Printf("error getting layer version: %s", err)
+		os.Exit(1)
+	}
+	log.Printf("Greatest version is %d. Bumping all versions...", greatestVersion)
+
+	// Elevate all regions to the greatest layer version found
+	err = balanceRegions(ctx, greatestVersion)
+	if err != nil {
+		cancel()
+		log.Printf("error balancing regions: %s", err)
+		os.Exit(1)
+	}
+
+	log.Printf("DONE! All layers should be version %d", greatestVersion)
+}