[release-1.11] Added hard reset conditions to RabbitMQ connections in the Adapter and Dispatcher (#1222)

knative-prow-robot · gabo1208 · web-flow · commit 060ecf0c23e8 · 2023-09-07T17:00:13.000Z
* implemented rabbitmq reconnection pod hard reset that kill the pods after 5 soft retries, now the state of the whole thing is clear to the user

* refactore to improve testing

---------

Co-authored-by: Gabriel Freites &lt;gfreites@vmware.com&gt;
diff --git a/cmd/dispatcher/main.go b/cmd/dispatcher/main.go
@@ -104,7 +104,7 @@ func main() {
 	}
 
 	var err error
-	rmqHelper := rabbit.NewRabbitMQConnectionHandler(1000, logger)
+	rmqHelper := rabbit.NewRabbitMQConnectionHandler(5, 1000, logger)
 	rmqHelper.Setup(ctx, rabbit.VHostHandler(
 		env.RabbitURL,
 		env.RabbitMQVhost),
diff --git a/cmd/ingress/main.go b/cmd/ingress/main.go
@@ -91,7 +91,7 @@ func main() {
 		logger.Errorw("failed to create the metrics exporter", zap.Error(err))
 	}
 
-	env.rmqHelper = rabbit.NewRabbitMQConnectionHandler(1000, logger)
+	env.rmqHelper = rabbit.NewRabbitMQConnectionHandler(5, 1000, logger)
 	env.rmqHelper.Setup(ctx, rabbit.VHostHandler(env.BrokerURL, env.RabbitMQVhost), rabbit.ChannelConfirm, rabbit.DialWrapper)
 
 	env.reporter = ingress.NewStatsReporter(env.ContainerName, kmeta.ChildName(env.PodName, uuid.New().String()))
diff --git a/pkg/adapter/adapter.go b/pkg/adapter/adapter.go
@@ -98,7 +98,7 @@ func (a *Adapter) start(stopCh <-chan struct{}) error {
 		zap.String("SinkURI", a.config.Sink))
 
 	if a.rmqHelper == nil {
-		a.rmqHelper = rabbit.NewRabbitMQConnectionHandler(1000, logger)
+		a.rmqHelper = rabbit.NewRabbitMQConnectionHandler(5, 1000, logger)
 		a.rmqHelper.Setup(a.context, rabbit.VHostHandler(a.config.RabbitURL, a.config.Vhost), rabbit.ChannelQoS, rabbit.DialWrapper)
 	}
 	return a.PollForMessages(stopCh)
diff --git a/pkg/adapter/adapter_test.go b/pkg/adapter/adapter_test.go
@@ -250,7 +250,7 @@ func TestAdapter_PollForMessages(t *testing.T) {
 		context:   ctx,
 		logger:    zap.NewNop(),
 		reporter:  statsReporter,
-		rmqHelper: rabbit.NewRabbitMQConnectionHandler(500, zap.NewNop().Sugar()),
+		rmqHelper: rabbit.NewRabbitMQConnectionHandler(5, 500, zap.NewNop().Sugar()),
 	}
 	a.rmqHelper.Setup(ctx, "", nil, rabbit.ValidDial)
 
diff --git a/pkg/rabbit/connections_handler.go b/pkg/rabbit/connections_handler.go
@@ -55,10 +55,11 @@ type RabbitMQChannelInterface interface {
 }
 
 type RabbitMQConnectionHandler struct {
-	firstSetup    bool
-	cycleDuration time.Duration
-	Connection    RabbitMQConnectionWrapperInterface
-	Channel       RabbitMQChannelInterface
+	firstSetup                              bool
+	reconnTries, reconnectionTriesThreshold int
+	cycleDuration                           time.Duration
+	Connection                              RabbitMQConnectionWrapperInterface
+	Channel                                 RabbitMQChannelInterface
 
 	logger *zap.SugaredLogger
 }
@@ -102,11 +103,12 @@ func (r *RabbitMQConnection) IsClosed() bool {
 	return true
 }
 
-func NewRabbitMQConnectionHandler(cycleDuration time.Duration, logger *zap.SugaredLogger) RabbitMQConnectionsHandlerInterface {
+func NewRabbitMQConnectionHandler(reconnectionTriesThreshold int, cycleDuration time.Duration, logger *zap.SugaredLogger) RabbitMQConnectionsHandlerInterface {
 	return &RabbitMQConnectionHandler{
-		firstSetup:    true,
-		cycleDuration: cycleDuration,
-		logger:        logger,
+		reconnectionTriesThreshold: reconnectionTriesThreshold,
+		firstSetup:                 true,
+		cycleDuration:              cycleDuration,
+		logger:                     logger,
 	}
 }
 
@@ -127,8 +129,7 @@ func (r *RabbitMQConnectionHandler) createConnectionAndChannel(
 	ctx context.Context,
 	rabbitMQURL string,
 	configFunction func(RabbitMQConnectionInterface, RabbitMQChannelInterface) error,
-	dialFunc func(string) (RabbitMQConnectionWrapperInterface, error)) {
-	var err error
+	dialFunc func(string) (RabbitMQConnectionWrapperInterface, error)) (err error) {
 	for {
 		// create the connection
 		r.Connection, err = r.createConnection(rabbitMQURL, dialFunc)
@@ -142,10 +143,17 @@ func (r *RabbitMQConnectionHandler) createConnectionAndChannel(
 		}
 		// setup completed successfully
 		if err == nil {
-			break
+			r.reconnTries = 0
+			return
 		}
 		r.logger.Error(err)
 		r.closeRabbitMQConnections()
+
+		r.reconnTries += 1
+		if r.reconnTries > r.reconnectionTriesThreshold {
+			err = errors.New("could not communicate to rabbitmq, restarting pods...")
+			return
+		}
 	}
 }
 
@@ -164,7 +172,9 @@ func (r *RabbitMQConnectionHandler) watchRabbitMQConnections(
 		case <-r.GetChannel().NotifyClose(make(chan *amqp091.Error)):
 		}
 		r.closeRabbitMQConnections()
-		r.createConnectionAndChannel(ctx, rabbitMQURL, configFunction, dialFunc)
+		if err := r.createConnectionAndChannel(ctx, rabbitMQURL, configFunction, dialFunc); err != nil {
+			r.logger.Fatal(err)
+		}
 	}
 }
 
diff --git a/pkg/rabbit/connections_handler_test.go b/pkg/rabbit/connections_handler_test.go
@@ -29,7 +29,7 @@ import (
 func Test_ValidSetupRabbitMQ(t *testing.T) {
 	logger := zap.NewNop().Sugar()
 	ctx, cancelFunc := context.WithCancel(context.Background())
-	rabbitMQHelper := NewRabbitMQConnectionHandler(100, logger).(*RabbitMQConnectionHandler)
+	rabbitMQHelper := NewRabbitMQConnectionHandler(1, 100, logger).(*RabbitMQConnectionHandler)
 	rabbitMQHelper.Setup(ctx, "amqp://localhost:5672/%2f", ConfigTest, ValidDial)
 	if rabbitMQHelper.Connection == nil || rabbitMQHelper.Channel == nil {
 		t.Errorf("rabbitMQHelper connection and channel should be set %s %s", rabbitMQHelper.Connection, rabbitMQHelper.Channel)
@@ -41,15 +41,15 @@ func Test_InvalidSetupRabbitMQ(t *testing.T) {
 	var err error
 	logger := zap.NewNop().Sugar()
 	// test invalid connection setup
-	rabbitMQHelper := NewRabbitMQConnectionHandler(100, logger).(*RabbitMQConnectionHandler)
-	rabbitMQHelper.Connection, err = rabbitMQHelper.createConnection("amqp://localhost:5672/%2f", BadConnectionDial)
+	rabbitMQHelper := NewRabbitMQConnectionHandler(1, 100, logger).(*RabbitMQConnectionHandler)
+	err = rabbitMQHelper.createConnectionAndChannel(context.TODO(), "amqp://localhost:5672/%2f", nil, BadConnectionDial)
 	if err == nil || rabbitMQHelper.GetConnection() != nil {
 		t.Errorf("unexpected error == nil when setting up invalid connection %s %s %s", rabbitMQHelper.GetConnection(), rabbitMQHelper.GetChannel(), err)
 	}
 	rabbitMQHelper.Connection = nil
 	rabbitMQHelper.closeRabbitMQConnections()
 	// test invalid channel setup
-	rabbitMQHelper = NewRabbitMQConnectionHandler(100, logger).(*RabbitMQConnectionHandler)
+	rabbitMQHelper = NewRabbitMQConnectionHandler(1, 100, logger).(*RabbitMQConnectionHandler)
 	rabbitMQHelper.Connection, _ = rabbitMQHelper.createConnection("amqp://localhost:5672/%2f", BadChannelDial)
 	rabbitMQHelper.Channel, err = rabbitMQHelper.createChannel()
 	if err == nil || rabbitMQHelper.GetChannel() != nil {
@@ -58,7 +58,7 @@ func Test_InvalidSetupRabbitMQ(t *testing.T) {
 	rabbitMQHelper.closeRabbitMQConnections()
 	rabbitMQHelper.Connection, rabbitMQHelper.Channel = nil, nil
 	// test invalid config setup
-	rabbitMQHelper = NewRabbitMQConnectionHandler(100, logger).(*RabbitMQConnectionHandler)
+	rabbitMQHelper = NewRabbitMQConnectionHandler(1, 100, logger).(*RabbitMQConnectionHandler)
 
 	rabbitMQHelper.Connection, _ = rabbitMQHelper.createConnection("amqp://localhost:5672/%2f", ValidDial)
 	rabbitMQHelper.Channel, _ = rabbitMQHelper.createChannel()
@@ -104,9 +104,8 @@ func Test_WatchConnectionsRabbitMQ(t *testing.T) {
 			ctx, cancelFunc := context.WithCancel(context.Background())
 			logger := zap.NewNop().Sugar()
 			// test invalid connection setup
-			rabbitMQHelper := NewRabbitMQConnectionHandler(100, logger).(*RabbitMQConnectionHandler)
-			rabbitMQHelper.Connection, _ = rabbitMQHelper.createConnection("amqp://localhost:5672/%2f", ValidDial)
-			rabbitMQHelper.Channel, _ = rabbitMQHelper.createChannel()
+			rabbitMQHelper := NewRabbitMQConnectionHandler(1, 100, logger).(*RabbitMQConnectionHandler)
+			rabbitMQHelper.createConnectionAndChannel(context.TODO(), "amqp://localhost:5672/%2f", nil, ValidDial)
 			go func() {
 				time.Sleep(time.Millisecond * 200)
 				if tt.endFunc == "connection" {

Original file line number	Diff line number	Diff line change
`@@ -104,7 +104,7 @@ func main() {`
`104`	`104`	`}`
`105`	`105`
`106`	`106`	`var err error`
`107`		`- rmqHelper := rabbit.NewRabbitMQConnectionHandler(1000, logger)`
	`107`	`+ rmqHelper := rabbit.NewRabbitMQConnectionHandler(5, 1000, logger)`
`108`	`108`	`rmqHelper.Setup(ctx, rabbit.VHostHandler(`
`109`	`109`	`env.RabbitURL,`
`110`	`110`	`env.RabbitMQVhost),`
Original file line number	Diff line number	Diff line change
`@@ -91,7 +91,7 @@ func main() {`
`91`	`91`	`logger.Errorw("failed to create the metrics exporter", zap.Error(err))`
`92`	`92`	`}`
`93`	`93`
`94`		`- env.rmqHelper = rabbit.NewRabbitMQConnectionHandler(1000, logger)`
	`94`	`+ env.rmqHelper = rabbit.NewRabbitMQConnectionHandler(5, 1000, logger)`
`95`	`95`	`env.rmqHelper.Setup(ctx, rabbit.VHostHandler(env.BrokerURL, env.RabbitMQVhost), rabbit.ChannelConfirm, rabbit.DialWrapper)`
`96`	`96`
`97`	`97`	`env.reporter = ingress.NewStatsReporter(env.ContainerName, kmeta.ChildName(env.PodName, uuid.New().String()))`
Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,7 @@ func (a *Adapter) start(stopCh <-chan struct{}) error {`
`98`	`98`	`zap.String("SinkURI", a.config.Sink))`
`99`	`99`
`100`	`100`	`if a.rmqHelper == nil {`
`101`		`- a.rmqHelper = rabbit.NewRabbitMQConnectionHandler(1000, logger)`
	`101`	`+ a.rmqHelper = rabbit.NewRabbitMQConnectionHandler(5, 1000, logger)`
`102`	`102`	`a.rmqHelper.Setup(a.context, rabbit.VHostHandler(a.config.RabbitURL, a.config.Vhost), rabbit.ChannelQoS, rabbit.DialWrapper)`
`103`	`103`	`}`
`104`	`104`	`return a.PollForMessages(stopCh)`
Original file line number	Diff line number	Diff line change
`@@ -250,7 +250,7 @@ func TestAdapter_PollForMessages(t *testing.T) {`
`250`	`250`	`context: ctx,`
`251`	`251`	`logger: zap.NewNop(),`
`252`	`252`	`reporter: statsReporter,`
`253`		`- rmqHelper: rabbit.NewRabbitMQConnectionHandler(500, zap.NewNop().Sugar()),`
	`253`	`+ rmqHelper: rabbit.NewRabbitMQConnectionHandler(5, 500, zap.NewNop().Sugar()),`
`254`	`254`	`}`
`255`	`255`	`a.rmqHelper.Setup(ctx, "", nil, rabbit.ValidDial)`
`256`	`256`
Original file line number	Diff line number	Diff line change
`@@ -55,10 +55,11 @@ type RabbitMQChannelInterface interface {`
`55`	`55`	`}`
`56`	`56`
`57`	`57`	`type RabbitMQConnectionHandler struct {`
`58`		`- firstSetup bool`
`59`		`- cycleDuration time.Duration`
`60`		`- Connection RabbitMQConnectionWrapperInterface`
`61`		`- Channel RabbitMQChannelInterface`
	`58`	`+ firstSetup bool`
	`59`	`+ reconnTries, reconnectionTriesThreshold int`
	`60`	`+ cycleDuration time.Duration`
	`61`	`+ Connection RabbitMQConnectionWrapperInterface`
	`62`	`+ Channel RabbitMQChannelInterface`
`62`	`63`
`63`	`64`	`logger *zap.SugaredLogger`
`64`	`65`	`}`
`@@ -102,11 +103,12 @@ func (r *RabbitMQConnection) IsClosed() bool {`
`102`	`103`	`return true`
`103`	`104`	`}`
`104`	`105`
`105`		`-func NewRabbitMQConnectionHandler(cycleDuration time.Duration, logger *zap.SugaredLogger) RabbitMQConnectionsHandlerInterface {`
	`106`	`+func NewRabbitMQConnectionHandler(reconnectionTriesThreshold int, cycleDuration time.Duration, logger *zap.SugaredLogger) RabbitMQConnectionsHandlerInterface {`
`106`	`107`	`return &RabbitMQConnectionHandler{`
`107`		`- firstSetup: true,`
`108`		`- cycleDuration: cycleDuration,`
`109`		`- logger: logger,`
	`108`	`+ reconnectionTriesThreshold: reconnectionTriesThreshold,`
	`109`	`+ firstSetup: true,`
	`110`	`+ cycleDuration: cycleDuration,`
	`111`	`+ logger: logger,`
`110`	`112`	`}`
`111`	`113`	`}`
`112`	`114`
`@@ -127,8 +129,7 @@ func (r *RabbitMQConnectionHandler) createConnectionAndChannel(`
`127`	`129`	`ctx context.Context,`
`128`	`130`	`rabbitMQURL string,`
`129`	`131`	`configFunction func(RabbitMQConnectionInterface, RabbitMQChannelInterface) error,`
`130`		`- dialFunc func(string) (RabbitMQConnectionWrapperInterface, error)) {`
`131`		`- var err error`
	`132`	`+ dialFunc func(string) (RabbitMQConnectionWrapperInterface, error)) (err error) {`
`132`	`133`	`for {`
`133`	`134`	`// create the connection`
`134`	`135`	`r.Connection, err = r.createConnection(rabbitMQURL, dialFunc)`
`@@ -142,10 +143,17 @@ func (r *RabbitMQConnectionHandler) createConnectionAndChannel(`
`142`	`143`	`}`
`143`	`144`	`// setup completed successfully`
`144`	`145`	`if err == nil {`
`145`		`- break`
	`146`	`+ r.reconnTries = 0`
	`147`	`+ return`
`146`	`148`	`}`
`147`	`149`	`r.logger.Error(err)`
`148`	`150`	`r.closeRabbitMQConnections()`
	`151`	`+`
	`152`	`+ r.reconnTries += 1`
	`153`	`+ if r.reconnTries > r.reconnectionTriesThreshold {`
	`154`	`+ err = errors.New("could not communicate to rabbitmq, restarting pods...")`
	`155`	`+ return`
	`156`	`+ }`
`149`	`157`	`}`
`150`	`158`	`}`
`151`	`159`
`@@ -164,7 +172,9 @@ func (r *RabbitMQConnectionHandler) watchRabbitMQConnections(`
`164`	`172`	`case <-r.GetChannel().NotifyClose(make(chan *amqp091.Error)):`
`165`	`173`	`}`
`166`	`174`	`r.closeRabbitMQConnections()`
`167`		`- r.createConnectionAndChannel(ctx, rabbitMQURL, configFunction, dialFunc)`
	`175`	`+ if err := r.createConnectionAndChannel(ctx, rabbitMQURL, configFunction, dialFunc); err != nil {`
	`176`	`+ r.logger.Fatal(err)`
	`177`	`+ }`
`168`	`178`	`}`
`169`	`179`	`}`
`170`	`180`