From ad9971879778c90155d82c0f5c29b62c375c44d6 Mon Sep 17 00:00:00 2001
From: Delnat Wito <wito.delnat@gmail.com>
Date: Tue, 22 Oct 2024 12:15:05 +0200
Subject: [PATCH] chore: add NATS retry

---
 go.mod                |  1 +
 go.sum                |  2 ++
 pkg/event/bus/nats.go | 12 +++++++++++-
 pkg/utils/consts.go   |  3 +++
 4 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index 7c228faa479..2aef9c43bf7 100644
--- a/go.mod
+++ b/go.mod
@@ -86,6 +86,7 @@ require (
 	github.com/agnivade/levenshtein v1.1.1 // indirect
 	github.com/alecthomas/chroma v0.10.0 // indirect
 	github.com/andybalholm/brotli v1.0.5 // indirect
+	github.com/avast/retry-go/v4 v4.6.0 // indirect
 	github.com/aymerick/douceur v0.2.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 // indirect
diff --git a/go.sum b/go.sum
index 671b4cd184a..2a1ba9f27aa 100644
--- a/go.sum
+++ b/go.sum
@@ -87,6 +87,8 @@ github.com/arbovm/levenshtein v0.0.0-20160628152529-48b4e1c0c4d0/go.mod h1:t2tdK
 github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio=
 github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs=
 github.com/atomicgo/cursor v0.0.1/go.mod h1:cBON2QmmrysudxNBFthvMtN32r3jxVRIvzkUiF/RuIk=
+github.com/avast/retry-go/v4 v4.6.0 h1:K9xNA+KeB8HHc2aWFuLb25Offp+0iVRXEvFx8IinRJA=
+github.com/avast/retry-go/v4 v4.6.0/go.mod h1:gvWlPhBVsvBbLkVGDg/KwvBv0bEkCOLRRSHKIr2PyOE=
 github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
 github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
diff --git a/pkg/event/bus/nats.go b/pkg/event/bus/nats.go
index 79e61d1ebff..c995606f7d0 100644
--- a/pkg/event/bus/nats.go
+++ b/pkg/event/bus/nats.go
@@ -6,11 +6,13 @@ import (
 	"sync"
 	"time"
 
+	"github.com/avast/retry-go/v4"
 	"github.com/nats-io/nats.go"
 
 	"github.com/kubeshop/testkube/pkg/api/v1/testkube"
 	"github.com/kubeshop/testkube/pkg/event/kind/common"
 	"github.com/kubeshop/testkube/pkg/log"
+	"github.com/kubeshop/testkube/pkg/utils"
 )
 
 var (
@@ -80,7 +82,15 @@ func NewNATSEncodedConnection(cfg ConnectionConfig, opts ...nats.Option) (*nats.
 func NewNATSConnection(cfg ConnectionConfig, opts ...nats.Option) (*nats.Conn, error) {
 	opts = append(opts, optsFromConfig(cfg)...)
 
-	nc, err := nats.Connect(cfg.NatsURI, opts...)
+	log.DefaultLogger.Infoln("Connecting to NATS")
+	nc, err := retry.DoWithData(
+		func() (*nats.Conn, error) {
+			return nats.Connect(cfg.NatsURI, opts...)
+		},
+		retry.DelayType(retry.FixedDelay),
+		retry.Delay(utils.DefaultRetryDelay),
+		retry.Attempts(20),
+	)
 	if err != nil {
 		log.DefaultLogger.Fatalw("error connecting to nats", "error", err)
 		return nil, err
diff --git a/pkg/utils/consts.go b/pkg/utils/consts.go
index 043bb8500c0..890c7cfac5b 100644
--- a/pkg/utils/consts.go
+++ b/pkg/utils/consts.go
@@ -1,6 +1,9 @@
 package utils
 
+import "time"
+
 const (
 	// DefaultDockerRegistry is the default registry used when no registry is specified in the image name.
 	DefaultDockerRegistry = "https://index.docker.io/v1/"
+	DefaultRetryDelay     = time.Second * 3
 )