Skip to content

Commit 708ae2d

Browse files
moshloopmoshloop
committed
chore: improve retry handling
1 parent 59bf0ae commit 708ae2d

5 files changed

Lines changed: 96 additions & 19 deletions

File tree

‎Makefile‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,10 +91,11 @@ $(LOCALBIN):
9191
$(ENVTEST): $(LOCALBIN)
9292
test -s $(LOCALBIN)/setup-envtest || GOBIN=$(LOCALBIN) go install sigs.k8s.io/controller-runtime/tools/setup-envtest@latest
9393

94+
95+
.PHONY: test
9496
test: $(ENVTEST)
9597
KUBEBUILDER_ASSETS="$(shell $(ENVTEST) use --bin-dir $(LOCALBIN) -p path)" go test -v ./...
9698

97-
9899
CONTROLLER_GEN = $(LOCALBIN)/controller-gen
99100

100101
.PHONY: controller-gen

‎go.mod‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,11 @@ require (
1010
github.com/aws/aws-sdk-go-v2/credentials v1.17.34
1111
github.com/aws/aws-sdk-go-v2/service/sns v1.31.3
1212
github.com/aws/aws-sdk-go-v2/service/sqs v1.34.3
13+
github.com/eko/gocache/lib/v4 v4.1.6
1314
github.com/flanksource/commons v1.36.1
1415
github.com/flanksource/duty v1.0.909
1516
github.com/flanksource/gomplate/v3 v3.24.56
17+
github.com/google/gops v0.3.28
1618
github.com/onsi/gomega v1.36.2
1719
github.com/samber/lo v1.47.0
1820
github.com/samber/oops v1.13.1
@@ -98,7 +100,6 @@ require (
98100
github.com/eapache/go-resiliency v1.7.0 // indirect
99101
github.com/eapache/go-xerial-snappy v0.0.0-20230731223053-c322873962e3 // indirect
100102
github.com/eapache/queue v1.1.0 // indirect
101-
github.com/eko/gocache/lib/v4 v4.1.6 // indirect
102103
github.com/eko/gocache/store/go_cache/v4 v4.2.2 // indirect
103104
github.com/emicklei/go-restful/v3 v3.12.1 // indirect
104105
github.com/emirpasic/gods v1.18.1 // indirect
@@ -144,7 +145,6 @@ require (
144145
github.com/google/gnostic-models v0.6.9-0.20230804172637-c7be7c783f49 // indirect
145146
github.com/google/go-cmp v0.6.0 // indirect
146147
github.com/google/gofuzz v1.2.0 // indirect
147-
github.com/google/gops v0.3.28 // indirect
148148
github.com/google/s2a-go v0.1.8 // indirect
149149
github.com/google/uuid v1.6.0 // indirect
150150
github.com/google/wire v0.6.0 // indirect

‎pkg/consumer.go‎

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import (
1212
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
1313
"sigs.k8s.io/yaml"
1414

15-
batchv1alpha1 "github.com/flanksource/batch-runner/pkg/apis/batch/v1"
15+
v1 "github.com/flanksource/batch-runner/pkg/apis/batch/v1"
1616
"github.com/flanksource/commons/logger"
1717
"github.com/flanksource/duty/context"
1818
"github.com/flanksource/duty/shell"
@@ -31,6 +31,8 @@ import (
3131
_ "gocloud.dev/pubsub/rabbitpubsub"
3232
)
3333

34+
var retry = NewRetryCache()
35+
3436
func pretty(o any) string {
3537
s, err := json.MarshalIndent(o, "", " ")
3638
if err != nil {
@@ -42,25 +44,23 @@ func pretty(o any) string {
4244
return err.Error()
4345
}
4446
return string(b)
45-
4647
}
47-
func RunConsumer(rootCtx context.Context, config *batchv1alpha1.Config) error {
48+
func RunConsumer(rootCtx context.Context, config *v1.Config) error {
4849
if config.LogLevel != "" {
4950
logger.StandardLogger().SetLogLevel(config.LogLevel)
5051
rootCtx.Infof("Set log level to %s => %v", config.LogLevel, rootCtx.Logger.GetLevel())
5152
}
5253

53-
rootCtx.Infof("Config: \n%+v", pretty(config))
54+
rootCtx.Tracef("Config: \n%+v", pretty(config))
5455

5556
sub, err := Subscribe(rootCtx, config)
5657
if err != nil {
5758
return oops.Wrapf(err, "Error building URL")
5859
}
5960

60-
rootCtx.Infof("Receiving messages from %+v", sub)
61+
rootCtx.Infof("Consuming from %s", config.String())
6162
ctx2, cancel := gocontext.WithCancel(gocontext.Background())
6263
shutdown.AddHook(func() {
63-
rootCtx.Infof("Shutting down consumer")
6464
cancel()
6565
})
6666

@@ -160,15 +160,27 @@ func RunConsumer(rootCtx context.Context, config *batchv1alpha1.Config) error {
160160
continue
161161
}
162162

163-
_delay := time.Second * 5
164-
if msg.Nackable() {
165-
msg.Nack()
163+
if exec.Retry == nil {
164+
exec.Retry = &v1.Retry{
165+
Attempts: 3,
166+
Delay: 30,
167+
}
166168
}
169+
167170
if err != nil {
168-
ctx.Errorf("Error running, (retrying in %s %v\n%s", _delay, err, exec.Script)
171+
ctx.Errorf("Error running %s: %s\n%s", exec.Script, err, details)
169172
} else {
170173
ctx.Errorf("Script returned non-zero exit code: %s", details)
171-
time.Sleep(_delay)
174+
}
175+
176+
delay := retry.GetBackoff(ctx, msg.LoggableID, exec.Retry)
177+
if delay != nil {
178+
if msg.Nackable() {
179+
msg.Nack()
180+
}
181+
time.Sleep(*delay)
182+
} else {
183+
msg.Ack()
172184
}
173185

174186
} else {

‎pkg/consumer_test.go‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,9 @@ import (
1313
"github.com/aws/aws-sdk-go-v2/service/sns"
1414
"github.com/aws/aws-sdk-go-v2/service/sqs"
1515
"github.com/aws/aws-sdk-go-v2/service/sqs/types"
16+
batchv1alpha1 "github.com/flanksource/batch-runner/pkg/apis/batch/v1"
1617
"github.com/flanksource/duty"
18+
dutyKubernetes "github.com/flanksource/duty/kubernetes"
1719
. "github.com/onsi/gomega"
1820
corev1 "k8s.io/api/core/v1"
1921
v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -118,20 +120,21 @@ func TestSNSToSQSIntegration(t *testing.T) {
118120
configData, err := os.ReadFile("../config-pod.yaml")
119121
Expect(err).To(BeNil())
120122

121-
var config Config
123+
var config batchv1alpha1.Config
122124
Expect(yaml.Unmarshal(configData, &config)).To(BeNil())
123125
config.SQS.QueueArn = queueArn
124126
config.SQS.AccessKey.ValueStatic = "test"
125127
config.SQS.SecretKey.ValueStatic = "test"
126128
config.SQS.Endpoint = endpoint
127129
config.SQS.WaitTime = 3
128130

129-
config.client = kubernetes.NewForConfigOrDie(restConfig)
131+
client := kubernetes.NewForConfigOrDie(restConfig)
130132

131133
ctx, cancel, err := duty.Start("batch-runner", duty.ClientOnly)
132134
Expect(err).To(BeNil())
135+
ctx = ctx.WithLocalKubernetes(dutyKubernetes.NewKubeClient(ctx.Logger, client, restConfig))
133136
defer cancel()
134-
go RunConsumer(ctx, config)
137+
go RunConsumer(ctx, &config)
135138
// Publish message to SNS
136139
testMessage := "{\"a\": \"b\"}"
137140
_, err = snsClient.Publish(context.TODO(), &sns.PublishInput{
@@ -141,7 +144,7 @@ func TestSNSToSQSIntegration(t *testing.T) {
141144
Expect(err).To(BeNil())
142145

143146
findPod := func() *corev1.Pod {
144-
if pod, e := config.client.CoreV1().Pods("default").Get(context.TODO(), "batch-b", v1.GetOptions{}); e == nil {
147+
if pod, e := client.CoreV1().Pods("default").Get(context.TODO(), "batch-b", v1.GetOptions{}); e == nil {
145148
return pod
146149
}
147150
return nil
@@ -155,7 +158,6 @@ func TestSNSToSQSIntegration(t *testing.T) {
155158
Expect(pod.Name).To(Equal("batch-b"))
156159

157160
// Cleanup
158-
159161
_, err = sqsClient.DeleteQueue(context.TODO(), &sqs.DeleteQueueInput{
160162
QueueUrl: &queueURL,
161163
})

‎pkg/retry.go‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package pkg
2+
3+
import (
4+
"time"
5+
6+
gocache "github.com/eko/gocache/lib/v4/cache"
7+
v1 "github.com/flanksource/batch-runner/pkg/apis/batch/v1"
8+
"github.com/flanksource/duty/cache"
9+
"github.com/flanksource/duty/context"
10+
)
11+
12+
type RetryCache struct {
13+
items gocache.CacheInterface[*RetryItem]
14+
}
15+
16+
type RetryItem struct {
17+
LastAttempt time.Time
18+
Count int
19+
}
20+
21+
func NewRetryCache() *RetryCache {
22+
return &RetryCache{
23+
items: cache.NewCache[*RetryItem]("retry-cache", 24*time.Hour),
24+
}
25+
}
26+
27+
func (rc *RetryCache) GetBackoff(ctx context.Context, messageID string, retry *v1.Retry) *time.Duration {
28+
if retry == nil {
29+
retry = &v1.Retry{
30+
Attempts: 3,
31+
Delay: 30,
32+
}
33+
}
34+
baseDelay := time.Duration(retry.Delay) * time.Second
35+
36+
item, _ := rc.items.Get(ctx, messageID)
37+
if item == nil {
38+
rc.items.Set(ctx, messageID, &RetryItem{
39+
LastAttempt: time.Now(),
40+
Count: 1,
41+
})
42+
ctx.Warnf("Retrying in %s (1 of %d)", baseDelay, retry.Attempts)
43+
return &baseDelay
44+
}
45+
46+
item.Count++
47+
item.LastAttempt = time.Now()
48+
49+
if item.Count > retry.Attempts {
50+
ctx.Errorf("Max retries exceeded (%d)", retry.Attempts)
51+
rc.items.Delete(ctx, messageID)
52+
return nil
53+
}
54+
55+
ctx.Warnf("Retrying in %s (%d of %d)", baseDelay, item.Count, retry.Attempts)
56+
57+
return &baseDelay
58+
}
59+
60+
func (rc *RetryCache) Remove(ctx context.Context, messageID string) {
61+
rc.items.Delete(ctx, messageID)
62+
}

0 commit comments

Comments
 (0)