Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(layers): add layer balancer script #1643

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 37 additions & 0 deletions layer/scripts/layer-balancer/README.md
Original file line number Diff line number Diff line change
@@ -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 :-)
24 changes: 24 additions & 0 deletions layer/scripts/layer-balancer/go.mod
Original file line number Diff line number Diff line change
@@ -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
)
37 changes: 37 additions & 0 deletions layer/scripts/layer-balancer/go.sum
Original file line number Diff line number Diff line change
@@ -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=
292 changes: 292 additions & 0 deletions layer/scripts/layer-balancer/main.go
Original file line number Diff line number Diff line change
@@ -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)
}