Skip to content

Commit e78681e

Browse files
mrledmrled
committed
Add retries
1 parent 42891c9 commit e78681e

6 files changed

Lines changed: 167 additions & 53 deletions

File tree

‎hedgerules/cmd/hedgerules/main.go‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ import (
1818

1919
var version = "dev"
2020

21+
const defaultMaxRetries = 10
22+
2123
type config struct {
2224
OutputDir string `toml:"output-dir"`
2325
Region string `toml:"region"`
@@ -26,6 +28,7 @@ type config struct {
2628
ViewerRequestName string `toml:"viewer-request-name"`
2729
ViewerResponseName string `toml:"viewer-response-name"`
2830
DebugHeaders bool `toml:"debug-headers"`
31+
MaxRetries int `toml:"max-retries"`
2932
}
3033

3134
func main() {
@@ -60,6 +63,7 @@ func runDeploy(args []string) {
6063
dryRun := fs.Bool("dry-run", false, "parse and validate only, print plan")
6164
region := fs.String("region", "", "AWS region override")
6265
debugHeaders := fs.Bool("debug-headers", false, "inject debug headers into viewer-response function")
66+
maxRetries := fs.Int("max-retries", -1, fmt.Sprintf("max AWS throttle retries (default %d, 0 disables retries)", defaultMaxRetries))
6367
fs.Parse(args)
6468

6569
// Resolve @FILE syntax for string flags
@@ -102,6 +106,12 @@ func runDeploy(args []string) {
102106
if *debugHeaders {
103107
cfg.DebugHeaders = true
104108
}
109+
if *maxRetries >= 0 {
110+
cfg.MaxRetries = *maxRetries
111+
} else if cfg.MaxRetries == 0 {
112+
// Neither flag nor config set: use default
113+
cfg.MaxRetries = defaultMaxRetries
114+
}
105115

106116
// Validate required config
107117
if cfg.OutputDir == "" {
@@ -205,52 +215,52 @@ func runDeploy(args []string) {
205215

206216
// Step 5: Resolve KVS ARNs
207217
fmt.Fprintf(os.Stderr, "Resolving KVS ARNs...\n")
208-
redirectsARN, err := functions.ResolveKVSARN(ctx, cfClient, cfg.RedirectsKVSName)
218+
redirectsARN, err := functions.ResolveKVSARN(ctx, cfClient, cfg.RedirectsKVSName, cfg.MaxRetries)
209219
if err != nil {
210220
fatal("resolving redirects KVS: %v", err)
211221
}
212222
fmt.Fprintf(os.Stderr, "Redirects KVS: %s\n", redirectsARN)
213223

214-
headersARN, err := functions.ResolveKVSARN(ctx, cfClient, cfg.HeadersKVSName)
224+
headersARN, err := functions.ResolveKVSARN(ctx, cfClient, cfg.HeadersKVSName, cfg.MaxRetries)
215225
if err != nil {
216226
fatal("resolving headers KVS: %v", err)
217227
}
218228
fmt.Fprintf(os.Stderr, "Headers KVS: %s\n", headersARN)
219229

220230
// Step 6: Sync redirects KVS
221231
fmt.Fprintf(os.Stderr, "Syncing redirects KVS...\n")
222-
existingRedirects, redirectEtag, err := kvs.FetchExistingKeys(ctx, kvsClient, redirectsARN)
232+
existingRedirects, redirectEtag, err := kvs.FetchExistingKeys(ctx, kvsClient, redirectsARN, cfg.MaxRetries)
223233
if err != nil {
224234
fatal("fetching existing redirects: %v", err)
225235
}
226236
redirectPlan := kvs.ComputeSyncPlan(redirectData, existingRedirects)
227237
fmt.Fprintf(os.Stderr, "Redirects: %d puts, %d deletes\n", len(redirectPlan.Puts), len(redirectPlan.Deletes))
228-
if err := kvs.Sync(ctx, kvsClient, redirectsARN, redirectEtag, redirectPlan); err != nil {
238+
if err := kvs.Sync(ctx, kvsClient, redirectsARN, redirectEtag, redirectPlan, cfg.MaxRetries); err != nil {
229239
fatal("syncing redirects: %v", err)
230240
}
231241

232242
// Step 7: Sync headers KVS
233243
fmt.Fprintf(os.Stderr, "Syncing headers KVS...\n")
234-
existingHeaders, headerEtag, err := kvs.FetchExistingKeys(ctx, kvsClient, headersARN)
244+
existingHeaders, headerEtag, err := kvs.FetchExistingKeys(ctx, kvsClient, headersARN, cfg.MaxRetries)
235245
if err != nil {
236246
fatal("fetching existing headers: %v", err)
237247
}
238248
headerPlan := kvs.ComputeSyncPlan(headerData, existingHeaders)
239249
fmt.Fprintf(os.Stderr, "Headers: %d puts, %d deletes\n", len(headerPlan.Puts), len(headerPlan.Deletes))
240-
if err := kvs.Sync(ctx, kvsClient, headersARN, headerEtag, headerPlan); err != nil {
250+
if err := kvs.Sync(ctx, kvsClient, headersARN, headerEtag, headerPlan, cfg.MaxRetries); err != nil {
241251
fatal("syncing headers: %v", err)
242252
}
243253

244254
// Step 8: Deploy CloudFront Functions
245255
fmt.Fprintf(os.Stderr, "Deploying viewer-request function...\n")
246256
requestCode := functions.BuildFunctionCode(functions.ViewerRequestJS, functions.KVSIDFromARN(redirectsARN), false)
247-
if err := functions.DeployFunction(ctx, cfClient, cfg.ViewerRequestName, requestCode, redirectsARN); err != nil {
257+
if err := functions.DeployFunction(ctx, cfClient, cfg.ViewerRequestName, requestCode, redirectsARN, cfg.MaxRetries); err != nil {
248258
fatal("deploying viewer-request function: %v", err)
249259
}
250260

251261
fmt.Fprintf(os.Stderr, "Deploying viewer-response function...\n")
252262
responseCode := functions.BuildFunctionCode(functions.ViewerResponseJS, functions.KVSIDFromARN(headersARN), cfg.DebugHeaders)
253-
if err := functions.DeployFunction(ctx, cfClient, cfg.ViewerResponseName, responseCode, headersARN); err != nil {
263+
if err := functions.DeployFunction(ctx, cfClient, cfg.ViewerResponseName, responseCode, headersARN, cfg.MaxRetries); err != nil {
254264
fatal("deploying viewer-response function: %v", err)
255265
}
256266

‎hedgerules/go.mod‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ require (
77
github.com/aws/aws-sdk-go-v2/config v1.28.7
88
github.com/aws/aws-sdk-go-v2/service/cloudfront v1.44.2
99
github.com/aws/aws-sdk-go-v2/service/cloudfrontkeyvaluestore v1.8.10
10+
github.com/aws/smithy-go v1.22.1
1011
)
1112

1213
require (
@@ -22,6 +23,6 @@ require (
2223
github.com/aws/aws-sdk-go-v2/service/sso v1.24.8 // indirect
2324
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.7 // indirect
2425
github.com/aws/aws-sdk-go-v2/service/sts v1.33.3 // indirect
25-
github.com/aws/smithy-go v1.22.1 // indirect
26+
github.com/aws/smithy-go v1.22.1
2627
github.com/jmespath/go-jmespath v0.4.0 // indirect
2728
)

‎hedgerules/internal/functions/deploy.go‎

Lines changed: 52 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88

99
"github.com/aws/aws-sdk-go-v2/service/cloudfront"
1010
cftypes "github.com/aws/aws-sdk-go-v2/service/cloudfront/types"
11+
"github.com/mrled/hedgerules/hedgerules/internal/retry"
1112
)
1213

1314
// KVSIDFromARN extracts the KVS ID (UUID) from a full KVS ARN.
@@ -33,11 +34,16 @@ type KVSARNResolver interface {
3334
}
3435

3536
// ResolveKVSARN resolves a KVS name to its ARN by listing all KVS and matching by name.
36-
func ResolveKVSARN(ctx context.Context, client KVSARNResolver, kvsName string) (string, error) {
37+
func ResolveKVSARN(ctx context.Context, client KVSARNResolver, kvsName string, maxRetries int) (string, error) {
3738
var marker *string
3839
for {
39-
resp, err := client.ListKeyValueStores(ctx, &cloudfront.ListKeyValueStoresInput{
40-
Marker: marker,
40+
var resp *cloudfront.ListKeyValueStoresOutput
41+
err := retry.Do(maxRetries, func() error {
42+
var e error
43+
resp, e = client.ListKeyValueStores(ctx, &cloudfront.ListKeyValueStoresInput{
44+
Marker: marker,
45+
})
46+
return e
4147
})
4248
if err != nil {
4349
return "", fmt.Errorf("listing key value stores: %w", err)
@@ -61,7 +67,7 @@ func ResolveKVSARN(ctx context.Context, client KVSARNResolver, kvsName string) (
6167

6268
// DeployFunction creates or updates a CloudFront Function with the given
6369
// JS source code and KVS association. It publishes the function to LIVE stage.
64-
func DeployFunction(ctx context.Context, client CFClient, name string, code []byte, kvsARN string) error {
70+
func DeployFunction(ctx context.Context, client CFClient, name string, code []byte, kvsARN string, maxRetries int) error {
6571
runtime := cftypes.FunctionRuntimeCloudfrontJs20
6672
kvAssoc := &cftypes.KeyValueStoreAssociations{
6773
Quantity: int32Ptr(1),
@@ -72,9 +78,14 @@ func DeployFunction(ctx context.Context, client CFClient, name string, code []by
7278

7379
// Check if function already exists
7480
var etag string
75-
descResp, err := client.DescribeFunction(ctx, &cloudfront.DescribeFunctionInput{
76-
Name: &name,
77-
Stage: cftypes.FunctionStageDevelopment,
81+
var descResp *cloudfront.DescribeFunctionOutput
82+
err := retry.Do(maxRetries, func() error {
83+
var e error
84+
descResp, e = client.DescribeFunction(ctx, &cloudfront.DescribeFunctionInput{
85+
Name: &name,
86+
Stage: cftypes.FunctionStageDevelopment,
87+
})
88+
return e
7889
})
7990

8091
var notFound *cftypes.NoSuchFunctionExists
@@ -85,30 +96,40 @@ func DeployFunction(ctx context.Context, client CFClient, name string, code []by
8596
if descResp != nil && descResp.ETag != nil {
8697
// Function exists, update it
8798
etag = *descResp.ETag
88-
updateResp, err := client.UpdateFunction(ctx, &cloudfront.UpdateFunctionInput{
89-
Name: &name,
90-
IfMatch: &etag,
91-
FunctionCode: code,
92-
FunctionConfig: &cftypes.FunctionConfig{
93-
Comment: strPtr(fmt.Sprintf("Managed by hedgerules: %s", name)),
94-
Runtime: runtime,
95-
KeyValueStoreAssociations: kvAssoc,
96-
},
99+
var updateResp *cloudfront.UpdateFunctionOutput
100+
err := retry.Do(maxRetries, func() error {
101+
var e error
102+
updateResp, e = client.UpdateFunction(ctx, &cloudfront.UpdateFunctionInput{
103+
Name: &name,
104+
IfMatch: &etag,
105+
FunctionCode: code,
106+
FunctionConfig: &cftypes.FunctionConfig{
107+
Comment: strPtr(fmt.Sprintf("Managed by hedgerules: %s", name)),
108+
Runtime: runtime,
109+
KeyValueStoreAssociations: kvAssoc,
110+
},
111+
})
112+
return e
97113
})
98114
if err != nil {
99115
return fmt.Errorf("updating function %s: %w", name, err)
100116
}
101117
etag = *updateResp.ETag
102118
} else {
103119
// Function doesn't exist, create it
104-
createResp, err := client.CreateFunction(ctx, &cloudfront.CreateFunctionInput{
105-
Name: &name,
106-
FunctionCode: code,
107-
FunctionConfig: &cftypes.FunctionConfig{
108-
Comment: strPtr(fmt.Sprintf("Managed by hedgerules: %s", name)),
109-
Runtime: runtime,
110-
KeyValueStoreAssociations: kvAssoc,
111-
},
120+
var createResp *cloudfront.CreateFunctionOutput
121+
err := retry.Do(maxRetries, func() error {
122+
var e error
123+
createResp, e = client.CreateFunction(ctx, &cloudfront.CreateFunctionInput{
124+
Name: &name,
125+
FunctionCode: code,
126+
FunctionConfig: &cftypes.FunctionConfig{
127+
Comment: strPtr(fmt.Sprintf("Managed by hedgerules: %s", name)),
128+
Runtime: runtime,
129+
KeyValueStoreAssociations: kvAssoc,
130+
},
131+
})
132+
return e
112133
})
113134
if err != nil {
114135
return fmt.Errorf("creating function %s: %w", name, err)
@@ -117,9 +138,13 @@ func DeployFunction(ctx context.Context, client CFClient, name string, code []by
117138
}
118139

119140
// Publish to LIVE stage
120-
_, err = client.PublishFunction(ctx, &cloudfront.PublishFunctionInput{
121-
Name: &name,
122-
IfMatch: &etag,
141+
err = retry.Do(maxRetries, func() error {
142+
var e error
143+
_, e = client.PublishFunction(ctx, &cloudfront.PublishFunctionInput{
144+
Name: &name,
145+
IfMatch: &etag,
146+
})
147+
return e
123148
})
124149
if err != nil {
125150
return fmt.Errorf("publishing function %s: %w", name, err)

‎hedgerules/internal/kvs/sync.go‎

Lines changed: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66

77
"github.com/aws/aws-sdk-go-v2/service/cloudfrontkeyvaluestore"
88
cfkvstypes "github.com/aws/aws-sdk-go-v2/service/cloudfrontkeyvaluestore/types"
9+
"github.com/mrled/hedgerules/hedgerules/internal/retry"
910
)
1011

1112
// KVSClient abstracts the CloudFront KeyValueStore API.
@@ -44,10 +45,15 @@ func ComputeSyncPlan(desired *Data, existingKeys map[string]string) *SyncPlan {
4445
}
4546

4647
// FetchExistingKeys retrieves all current keys and values from a KVS.
47-
func FetchExistingKeys(ctx context.Context, client KVSClient, kvsARN string) (map[string]string, string, error) {
48+
func FetchExistingKeys(ctx context.Context, client KVSClient, kvsARN string, maxRetries int) (map[string]string, string, error) {
4849
// Get ETag
49-
desc, err := client.DescribeKeyValueStore(ctx, &cloudfrontkeyvaluestore.DescribeKeyValueStoreInput{
50-
KvsARN: &kvsARN,
50+
var desc *cloudfrontkeyvaluestore.DescribeKeyValueStoreOutput
51+
err := retry.Do(maxRetries, func() error {
52+
var e error
53+
desc, e = client.DescribeKeyValueStore(ctx, &cloudfrontkeyvaluestore.DescribeKeyValueStoreInput{
54+
KvsARN: &kvsARN,
55+
})
56+
return e
5157
})
5258
if err != nil {
5359
return nil, "", fmt.Errorf("describing KVS: %w", err)
@@ -57,9 +63,14 @@ func FetchExistingKeys(ctx context.Context, client KVSClient, kvsARN string) (ma
5763
existing := make(map[string]string)
5864
var nextToken *string
5965
for {
60-
resp, err := client.ListKeys(ctx, &cloudfrontkeyvaluestore.ListKeysInput{
61-
KvsARN: &kvsARN,
62-
NextToken: nextToken,
66+
var resp *cloudfrontkeyvaluestore.ListKeysOutput
67+
err := retry.Do(maxRetries, func() error {
68+
var e error
69+
resp, e = client.ListKeys(ctx, &cloudfrontkeyvaluestore.ListKeysInput{
70+
KvsARN: &kvsARN,
71+
NextToken: nextToken,
72+
})
73+
return e
6374
})
6475
if err != nil {
6576
return nil, "", fmt.Errorf("listing KVS keys: %w", err)
@@ -82,7 +93,7 @@ const maxKeysPerBatch = 50
8293

8394
// Sync applies a SyncPlan to a CloudFront KVS using the batch UpdateKeys API.
8495
// Large operations are automatically split into multiple batches to respect AWS limits.
85-
func Sync(ctx context.Context, client KVSClient, kvsARN string, etag string, plan *SyncPlan) error {
96+
func Sync(ctx context.Context, client KVSClient, kvsARN string, etag string, plan *SyncPlan, maxRetries int) error {
8697
if len(plan.Puts) == 0 && len(plan.Deletes) == 0 {
8798
return nil
8899
}
@@ -134,11 +145,16 @@ func Sync(ctx context.Context, client KVSClient, kvsARN string, etag string, pla
134145
}
135146

136147
// Execute batch
137-
resp, err := client.UpdateKeys(ctx, &cloudfrontkeyvaluestore.UpdateKeysInput{
138-
KvsARN: &kvsARN,
139-
IfMatch: &currentETag,
140-
Puts: batchPuts,
141-
Deletes: batchDeletes,
148+
var resp *cloudfrontkeyvaluestore.UpdateKeysOutput
149+
err := retry.Do(maxRetries, func() error {
150+
var e error
151+
resp, e = client.UpdateKeys(ctx, &cloudfrontkeyvaluestore.UpdateKeysInput{
152+
KvsARN: &kvsARN,
153+
IfMatch: &currentETag,
154+
Puts: batchPuts,
155+
Deletes: batchDeletes,
156+
})
157+
return e
142158
})
143159
if err != nil {
144160
return fmt.Errorf("updating KVS keys (batch %d/%d puts, %d deletes): %w",

‎hedgerules/internal/kvs/sync_test.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,7 @@ func TestSync_SmallPlan(t *testing.T) {
170170
Deletes: []string{"/old"},
171171
}
172172

173-
err := Sync(context.Background(), mock, "arn:test", "etag-0", plan)
173+
err := Sync(context.Background(), mock, "arn:test", "etag-0", plan, 0)
174174
if err != nil {
175175
t.Fatalf("Sync failed: %v", err)
176176
}
@@ -202,7 +202,7 @@ func TestSync_ExactlyMaxBatch(t *testing.T) {
202202

203203
plan := &SyncPlan{Puts: puts}
204204

205-
err := Sync(context.Background(), mock, "arn:test", "etag-0", plan)
205+
err := Sync(context.Background(), mock, "arn:test", "etag-0", plan, 0)
206206
if err != nil {
207207
t.Fatalf("Sync failed: %v", err)
208208
}
@@ -228,7 +228,7 @@ func TestSync_MultipleBatches(t *testing.T) {
228228

229229
plan := &SyncPlan{Puts: puts}
230230

231-
err := Sync(context.Background(), mock, "arn:test", "etag-0", plan)
231+
err := Sync(context.Background(), mock, "arn:test", "etag-0", plan, 0)
232232
if err != nil {
233233
t.Fatalf("Sync failed: %v", err)
234234
}
@@ -281,7 +281,7 @@ func TestSync_MixedPutsAndDeletes(t *testing.T) {
281281
Deletes: deletes,
282282
}
283283

284-
err := Sync(context.Background(), mock, "arn:test", "etag-0", plan)
284+
err := Sync(context.Background(), mock, "arn:test", "etag-0", plan, 0)
285285
if err != nil {
286286
t.Fatalf("Sync failed: %v", err)
287287
}
@@ -313,7 +313,7 @@ func TestSync_EmptyPlan(t *testing.T) {
313313
mock := &mockKVSClient{nextETag: 0}
314314
plan := &SyncPlan{}
315315

316-
err := Sync(context.Background(), mock, "arn:test", "etag-0", plan)
316+
err := Sync(context.Background(), mock, "arn:test", "etag-0", plan, 0)
317317
if err != nil {
318318
t.Fatalf("Sync failed: %v", err)
319319
}

0 commit comments

Comments
 (0)