Skip to content

Commit 557a700

Browse files
feat(layers): add layer balancer script (#1643)
Co-authored-by: Leandro Damascena <[email protected]>
1 parent 83e216a commit 557a700

File tree

4 files changed

+390
-0
lines changed

4 files changed

+390
-0
lines changed
+37
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
<!-- markdownlint-disable MD041 MD043 -->
2+
# Layer balancer
3+
4+
This folder contains a Go project that balances the layer version of Lambda Powertools across all regions, so
5+
every region has the same layer version.
6+
7+
Before:
8+
9+
```text
10+
arn:aws:lambda:eu-central-1:017000801446:layer:AWSLambdaPowertoolsPythonV2:11
11+
...
12+
arn:aws:lambda:eu-west-1:017000801446:layer:AWSLambdaPowertoolsPythonV2:9
13+
```
14+
15+
After:
16+
17+
```text
18+
arn:aws:lambda:eu-central-1:017000801446:layer:AWSLambdaPowertoolsPythonV2:11
19+
...
20+
arn:aws:lambda:eu-west-1:017000801446:layer:AWSLambdaPowertoolsPythonV2:11
21+
```
22+
23+
## What's happening under the hood?
24+
25+
1. Query all regions to find the greatest version number
26+
2. Download the latest layer from eu-central-1
27+
3. Use the layer contents to bump the version on each region until it matches 1
28+
29+
## Requirements
30+
31+
* go >= 1.18
32+
33+
## How to use
34+
35+
1. Set your AWS_PROFILE to the correct profile
36+
2. `go run .`
37+
3. Profit :-)

layer/scripts/layer-balancer/go.mod

+24
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
module layerbalancer
2+
3+
go 1.18
4+
5+
require (
6+
github.com/aws/aws-sdk-go-v2 v1.16.16
7+
github.com/aws/aws-sdk-go-v2/config v1.17.8
8+
github.com/aws/aws-sdk-go-v2/service/lambda v1.24.6
9+
golang.org/x/sync v0.1.0
10+
)
11+
12+
require (
13+
github.com/aws/aws-sdk-go-v2/credentials v1.12.21 // indirect
14+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.17 // indirect
15+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.23 // indirect
16+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.17 // indirect
17+
github.com/aws/aws-sdk-go-v2/internal/ini v1.3.24 // indirect
18+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.17 // indirect
19+
github.com/aws/aws-sdk-go-v2/service/sso v1.11.23 // indirect
20+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.13.6 // indirect
21+
github.com/aws/aws-sdk-go-v2/service/sts v1.16.19 // indirect
22+
github.com/aws/smithy-go v1.13.3 // indirect
23+
github.com/jmespath/go-jmespath v0.4.0 // indirect
24+
)

layer/scripts/layer-balancer/go.sum

+37
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
github.com/aws/aws-sdk-go-v2 v1.16.16 h1:M1fj4FE2lB4NzRb9Y0xdWsn2P0+2UHVxwKyOa4YJNjk=
2+
github.com/aws/aws-sdk-go-v2 v1.16.16/go.mod h1:SwiyXi/1zTUZ6KIAmLK5V5ll8SiURNUYOqTerZPaF9k=
3+
github.com/aws/aws-sdk-go-v2/config v1.17.8 h1:b9LGqNnOdg9vR4Q43tBTVWk4J6F+W774MSchvKJsqnE=
4+
github.com/aws/aws-sdk-go-v2/config v1.17.8/go.mod h1:UkCI3kb0sCdvtjiXYiU4Zx5h07BOpgBTtkPu/49r+kA=
5+
github.com/aws/aws-sdk-go-v2/credentials v1.12.21 h1:4tjlyCD0hRGNQivh5dN8hbP30qQhMLBE/FgQR1vHHWM=
6+
github.com/aws/aws-sdk-go-v2/credentials v1.12.21/go.mod h1:O+4XyAt4e+oBAoIwNUYkRg3CVMscaIJdmZBOcPgJ8D8=
7+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.17 h1:r08j4sbZu/RVi+BNxkBJwPMUYY3P8mgSDuKkZ/ZN1lE=
8+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.17/go.mod h1:yIkQcCDYNsZfXpd5UX2Cy+sWA1jPgIhGTw9cOBzfVnQ=
9+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.23 h1:s4g/wnzMf+qepSNgTvaQQHNxyMLKSawNhKCPNy++2xY=
10+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.23/go.mod h1:2DFxAQ9pfIRy0imBCJv+vZ2X6RKxves6fbnEuSry6b4=
11+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.17 h1:/K482T5A3623WJgWT8w1yRAFK4RzGzEl7y39yhtn9eA=
12+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.17/go.mod h1:pRwaTYCJemADaqCbUAxltMoHKata7hmB5PjEXeu0kfg=
13+
github.com/aws/aws-sdk-go-v2/internal/ini v1.3.24 h1:wj5Rwc05hvUSvKuOF29IYb9QrCLjU+rHAy/x/o0DK2c=
14+
github.com/aws/aws-sdk-go-v2/internal/ini v1.3.24/go.mod h1:jULHjqqjDlbyTa7pfM7WICATnOv+iOhjletM3N0Xbu8=
15+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.17 h1:Jrd/oMh0PKQc6+BowB+pLEwLIgaQF29eYbe7E1Av9Ug=
16+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.17/go.mod h1:4nYOrY41Lrbk2170/BGkcJKBhws9Pfn8MG3aGqjjeFI=
17+
github.com/aws/aws-sdk-go-v2/service/lambda v1.24.6 h1:N7RkXX2SJbN+TCp295J3LdMR0KRFd2Bhi5nIO+svLQY=
18+
github.com/aws/aws-sdk-go-v2/service/lambda v1.24.6/go.mod h1:oTJIIluTaJCRT6xP1AZpuU3JwRHBC0Q5O4Hg+SUxFHw=
19+
github.com/aws/aws-sdk-go-v2/service/sso v1.11.23 h1:pwvCchFUEnlceKIgPUouBJwK81aCkQ8UDMORfeFtW10=
20+
github.com/aws/aws-sdk-go-v2/service/sso v1.11.23/go.mod h1:/w0eg9IhFGjGyyncHIQrXtU8wvNsTJOP0R6PPj0wf80=
21+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.13.6 h1:OwhhKc1P9ElfWbMKPIbMMZBV6hzJlL2JKD76wNNVzgQ=
22+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.13.6/go.mod h1:csZuQY65DAdFBt1oIjO5hhBR49kQqop4+lcuCjf2arA=
23+
github.com/aws/aws-sdk-go-v2/service/sts v1.16.19 h1:9pPi0PsFNAGILFfPCk8Y0iyEBGc6lu6OQ97U7hmdesg=
24+
github.com/aws/aws-sdk-go-v2/service/sts v1.16.19/go.mod h1:h4J3oPZQbxLhzGnk+j9dfYHi5qIOVJ5kczZd658/ydM=
25+
github.com/aws/smithy-go v1.13.3 h1:l7LYxGuzK6/K+NzJ2mC+VvLUbae0sL3bXU//04MkmnA=
26+
github.com/aws/smithy-go v1.13.3/go.mod h1:Tg+OJXh4MB2R/uN61Ko2f6hTZwB/ZYGOtib8J3gBHzA=
27+
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
28+
github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
29+
github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
30+
github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
31+
github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
32+
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
33+
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
34+
golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
35+
golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
36+
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
37+
gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

layer/scripts/layer-balancer/main.go

+292
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,292 @@
1+
package main
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"io/ioutil"
7+
"log"
8+
"net/http"
9+
"os"
10+
"os/signal"
11+
"sort"
12+
"sync"
13+
14+
"github.com/aws/aws-sdk-go-v2/aws"
15+
"github.com/aws/aws-sdk-go-v2/config"
16+
"github.com/aws/aws-sdk-go-v2/service/lambda"
17+
"github.com/aws/aws-sdk-go-v2/service/lambda/types"
18+
"golang.org/x/sync/errgroup"
19+
)
20+
21+
type LayerInfo struct {
22+
Name string
23+
Description string
24+
Architecture types.Architecture
25+
26+
LayerContentOnce sync.Once
27+
LayerContent []byte
28+
}
29+
30+
// canonicalLayers are the layers that we want to keep in sync across all regions
31+
var canonicalLayers = []LayerInfo{
32+
{
33+
Name: "AWSLambdaPowertoolsPythonV2",
34+
Description: "Lambda Powertools for Python [x86_64] with extra dependencies version bump",
35+
Architecture: types.ArchitectureX8664,
36+
},
37+
{
38+
Name: "AWSLambdaPowertoolsPythonV2-Arm64",
39+
Description: "Lambda Powertools for Python [arm64] with extra dependencies version bump",
40+
Architecture: types.ArchitectureArm64,
41+
},
42+
}
43+
44+
// regions are the regions that we want to keep in sync
45+
var regions = []string{
46+
"af-south-1",
47+
"eu-central-1",
48+
"us-east-1",
49+
"us-east-2",
50+
"us-west-1",
51+
"us-west-2",
52+
"ap-east-1",
53+
"ap-south-1",
54+
"ap-northeast-1",
55+
"ap-northeast-2",
56+
"ap-southeast-1",
57+
"ap-southeast-2",
58+
"ca-central-1",
59+
"eu-west-1",
60+
"eu-west-2",
61+
"eu-west-3",
62+
"eu-south-1",
63+
"eu-north-1",
64+
"sa-east-1",
65+
"ap-southeast-3",
66+
"ap-northeast-3",
67+
"me-south-1",
68+
}
69+
70+
// getLayerVersion returns the latest version of a layer in a region
71+
func getLayerVersion(ctx context.Context, layerName string, region string) (int64, error) {
72+
cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
73+
if err != nil {
74+
return 0, err
75+
}
76+
77+
lambdaSvc := lambda.NewFromConfig(cfg)
78+
79+
layerVersionsResult, err := lambdaSvc.ListLayerVersions(ctx, &lambda.ListLayerVersionsInput{
80+
LayerName: aws.String(layerName),
81+
MaxItems: aws.Int32(1),
82+
})
83+
if err != nil {
84+
return 0, err
85+
}
86+
87+
if len(layerVersionsResult.LayerVersions) == 0 {
88+
return 0, fmt.Errorf("no layer meets the search criteria %s - %s", layerName, region)
89+
}
90+
return layerVersionsResult.LayerVersions[0].Version, nil
91+
}
92+
93+
// getGreatestVersion returns the greatest version of a layer across all regions
94+
func getGreatestVersion(ctx context.Context) (int64, error) {
95+
var versions []int64
96+
97+
g, ctx := errgroup.WithContext(ctx)
98+
99+
for idx := range canonicalLayers {
100+
layer := &canonicalLayers[idx]
101+
102+
for _, region := range regions {
103+
layerName := layer.Name
104+
ctx := ctx
105+
region := region
106+
107+
g.Go(func() error {
108+
version, err := getLayerVersion(ctx, layerName, region)
109+
if err != nil {
110+
return err
111+
}
112+
113+
log.Printf("[%s] %s -> %d", layerName, region, version)
114+
115+
versions = append(versions, version)
116+
return nil
117+
})
118+
}
119+
}
120+
121+
if err := g.Wait(); err != nil {
122+
return 0, err
123+
}
124+
125+
// Find the maximum version by reverse sorting the versions array
126+
sort.Slice(versions, func(i, j int) bool { return versions[i] > versions[j] })
127+
return versions[0], nil
128+
}
129+
130+
// balanceRegionToVersion creates a new layer version in a region with the same contents as the canonical layer, until it matches the maxVersion
131+
func balanceRegionToVersion(ctx context.Context, region string, layer *LayerInfo, maxVersion int64) error {
132+
currentLayerVersion, err := getLayerVersion(ctx, layer.Name, region)
133+
if err != nil {
134+
return fmt.Errorf("error getting layer version: %w", err)
135+
}
136+
137+
cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
138+
if err != nil {
139+
return err
140+
}
141+
142+
lambdaSvc := lambda.NewFromConfig(cfg)
143+
144+
for i := currentLayerVersion; i < maxVersion; i++ {
145+
log.Printf("[%s] Bumping %s to version %d (max %d)", layer.Name, region, i, maxVersion)
146+
147+
payload, err := downloadCanonicalLayerZip(ctx, layer)
148+
if err != nil {
149+
return fmt.Errorf("error downloading canonical zip: %w", err)
150+
}
151+
152+
layerVersionResponse, err := lambdaSvc.PublishLayerVersion(ctx, &lambda.PublishLayerVersionInput{
153+
Content: &types.LayerVersionContentInput{
154+
ZipFile: payload,
155+
},
156+
LayerName: aws.String(layer.Name),
157+
CompatibleArchitectures: []types.Architecture{layer.Architecture},
158+
CompatibleRuntimes: []types.Runtime{types.RuntimePython37, types.RuntimePython38, types.RuntimePython39},
159+
Description: aws.String(layer.Description),
160+
LicenseInfo: aws.String("MIT-0"),
161+
})
162+
if err != nil {
163+
return fmt.Errorf("error publishing layer version: %w", err)
164+
}
165+
166+
_, err = lambdaSvc.AddLayerVersionPermission(ctx, &lambda.AddLayerVersionPermissionInput{
167+
Action: aws.String("lambda:GetLayerVersion"),
168+
LayerName: aws.String(layer.Name),
169+
Principal: aws.String("*"),
170+
StatementId: aws.String("PublicLayerAccess"),
171+
VersionNumber: layerVersionResponse.Version,
172+
})
173+
if err != nil {
174+
return fmt.Errorf("error making layer public: %w", err)
175+
}
176+
}
177+
178+
return nil
179+
}
180+
181+
// balanceRegions creates new layer versions in all regions with the same contents as the canonical layer, until they match the maxVersion
182+
func balanceRegions(ctx context.Context, maxVersion int64) error {
183+
g, ctx := errgroup.WithContext(ctx)
184+
185+
for idx := range canonicalLayers {
186+
layer := &canonicalLayers[idx]
187+
188+
for _, region := range regions {
189+
ctx := ctx
190+
region := region
191+
layer := layer
192+
version := maxVersion
193+
194+
g.Go(func() error {
195+
return balanceRegionToVersion(ctx, region, layer, version)
196+
})
197+
}
198+
}
199+
200+
if err := g.Wait(); err != nil {
201+
return err
202+
}
203+
204+
return nil
205+
}
206+
207+
// downloadCanonicalLayerZip downloads the canonical layer zip file that will be used to bump the versions later
208+
func downloadCanonicalLayerZip(ctx context.Context, layer *LayerInfo) ([]byte, error) {
209+
var innerErr error
210+
211+
layer.LayerContentOnce.Do(func() {
212+
// We use eu-central-1 as the canonical region to download the Layer from
213+
cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion("eu-central-1"))
214+
if err != nil {
215+
innerErr = err
216+
}
217+
218+
lambdaSvc := lambda.NewFromConfig(cfg)
219+
220+
// Gets the latest version of the layer
221+
version, err := getLayerVersion(ctx, layer.Name, "eu-central-1")
222+
if err != nil {
223+
innerErr = fmt.Errorf("error getting eu-central-1 layer version: %w", err)
224+
}
225+
226+
// Gets the Layer content URL from S3
227+
getLayerVersionResult, err := lambdaSvc.GetLayerVersion(ctx, &lambda.GetLayerVersionInput{
228+
LayerName: aws.String(layer.Name),
229+
VersionNumber: version,
230+
})
231+
if err != nil {
232+
innerErr = fmt.Errorf("error getting eu-central-1 layer download URL: %w", err)
233+
}
234+
235+
s3LayerUrl := getLayerVersionResult.Content.Location
236+
log.Printf("[%s] Downloading Layer from %s", layer.Name, *s3LayerUrl)
237+
238+
resp, err := http.Get(*s3LayerUrl)
239+
if err != nil {
240+
innerErr = err
241+
}
242+
defer resp.Body.Close()
243+
244+
body, err := ioutil.ReadAll(resp.Body)
245+
if err != nil {
246+
innerErr = err
247+
}
248+
249+
layer.LayerContent = body
250+
})
251+
252+
return layer.LayerContent, innerErr
253+
}
254+
255+
func main() {
256+
ctx := context.Background()
257+
258+
// Cancel everything if interrupted
259+
ctx, cancel := context.WithCancel(ctx)
260+
c := make(chan os.Signal, 1)
261+
signal.Notify(c, os.Interrupt)
262+
defer func() {
263+
signal.Stop(c)
264+
cancel()
265+
}()
266+
go func() {
267+
select {
268+
case <-c:
269+
cancel()
270+
case <-ctx.Done():
271+
}
272+
}()
273+
274+
// Find the greatest layer version across all regions
275+
greatestVersion, err := getGreatestVersion(ctx)
276+
if err != nil {
277+
cancel()
278+
log.Printf("error getting layer version: %s", err)
279+
os.Exit(1)
280+
}
281+
log.Printf("Greatest version is %d. Bumping all versions...", greatestVersion)
282+
283+
// Elevate all regions to the greatest layer version found
284+
err = balanceRegions(ctx, greatestVersion)
285+
if err != nil {
286+
cancel()
287+
log.Printf("error balancing regions: %s", err)
288+
os.Exit(1)
289+
}
290+
291+
log.Printf("DONE! All layers should be version %d", greatestVersion)
292+
}

0 commit comments

Comments
 (0)