fix: Use WaitGroup to start eventsoures and sensors (#819)

Makefile

@@ -148,7 +148,6 @@ codegen:
 	./hack/update-openapigen.sh
 	$(MAKE) swagger
 	./hack/update-api-docs.sh
-	./hack/update-mocks.sh
 	rm -rf ./vendor
 	go mod tidy
 	$(MAKE) manifests

eventbus/driver/nats.go

@@ -80,7 +80,7 @@ func (n *natsStreaming) Connect() (Connection, error) {
 		nats.NoReconnect(),
 		nats.DisconnectErrHandler(func(nc *nats.Conn, err error) {
 			conn.natsConnected = false
-			log.Error("NATS connection los", zap.Error(err))
+			log.Error("NATS connection lost", zap.Error(err))
 		}),
 		nats.ReconnectHandler(func(nnc *nats.Conn) {
 			conn.natsConnected = true

eventsources/eventing.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/argoproj/argo-events/common/logging"
@@ -249,9 +250,7 @@ func (e *EventSourceAdaptor) Start(ctx context.Context, stopCh <-chan struct{})
 	logger := logging.FromContext(ctx).Desugar()
 	logger.Info("Starting event source server...")
 	servers := GetEventingServers(e.eventSource)
-	cctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-	driver, err := eventbus.GetDriver(cctx, *e.eventBusConfig, e.eventBusSubject, e.hostname)
+	driver, err := eventbus.GetDriver(ctx, *e.eventBusConfig, e.eventBusSubject, e.hostname)
 	if err != nil {
 		logger.Error("failed to get eventbus driver", zap.Error(err))
 		return err
@@ -263,8 +262,13 @@ func (e *EventSourceAdaptor) Start(ctx context.Context, stopCh <-chan struct{})
 	}
 	defer e.eventBusConn.Close()
 
+	cctx, cancel := context.WithCancel(ctx)
+	wg := &sync.WaitGroup{}
+
 	// Daemon to reconnect
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		logger.Info("starting eventbus connection daemon...")
 		ticker := time.NewTicker(5 * time.Second)
 		for {
@@ -296,7 +300,9 @@ func (e *EventSourceAdaptor) Start(ctx context.Context, stopCh <-chan struct{})
 				// Continue starting other event services instead of failing all of them
 				continue
 			}
+			wg.Add(1)
 			go func(s EventingServer) {
+				defer wg.Done()
 				err := s.StartListening(cctx, func(data []byte) error {
 					event := cloudevents.NewEvent()
 					event.SetID(fmt.Sprintf("%x", uuid.New()))
@@ -327,5 +333,7 @@ func (e *EventSourceAdaptor) Start(ctx context.Context, stopCh <-chan struct{})
 	logger.Info("Eventing server started.")
 	<-stopCh
 	logger.Info("Shutting down...")
+	cancel()
+	wg.Wait()
 	return nil
 }

sensors/listener.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/Knetic/govaluate"
@@ -65,9 +66,11 @@ func (sensorCtx *SensorContext) ListenEvents(ctx context.Context, stopCh <-chan
 	}
 
 	cctx, cancel := context.WithCancel(ctx)
-	defer cancel()
+	wg := &sync.WaitGroup{}
 	for k, v := range triggerMapping {
+		wg.Add(1)
 		go func(depExpression string, triggers []v1alpha1.Trigger) {
+			defer wg.Done()
 			// Calculate dependencies of each group of triggers.
 			de := strings.ReplaceAll(depExpression, "-", "\\-")
 			expr, err := govaluate.NewEvaluableExpression(de)
@@ -139,8 +142,11 @@ func (sensorCtx *SensorContext) ListenEvents(ctx context.Context, stopCh <-chan
 				}
 			}
 
+			wg1 := &sync.WaitGroup{}
 			closeSubCh := make(chan struct{})
+			wg1.Add(1)
 			go func() {
+				defer wg1.Done()
 				logger.Sugar().Infof("started to subscribe events for triggers %s with client %s", fmt.Sprintf("[%s]", strings.Join(triggerNames, " ")), clientID)
 				err = ebDriver.SubscribeEventSources(cctx, conn, closeSubCh, depExpression, deps, filterFunc, actionFunc)
 				if err != nil {
@@ -156,6 +162,7 @@ func (sensorCtx *SensorContext) ListenEvents(ctx context.Context, stopCh <-chan
 				case <-cctx.Done():
 					logger.Sugar().Infof("exiting eventbus connection daemon for client %s...", clientID)
 					ticker.Stop()
+					wg1.Wait()
 					return
 				case <-ticker.C:
 					if conn == nil || conn.IsClosed() {
@@ -168,7 +175,9 @@ func (sensorCtx *SensorContext) ListenEvents(ctx context.Context, stopCh <-chan
 						logger.Info("reconnected to NATS streaming server.", zap.Any("clientID", clientID))
 						closeSubCh <- struct{}{}
 						time.Sleep(2 * time.Second)
+						wg1.Add(1)
 						go func() {
+							defer wg1.Done()
 							logger.Sugar().Infof("started to re-subscribe events for triggers %s with client %s", fmt.Sprintf("[%s]", strings.Join(triggerNames, " ")), clientID)
 							err = ebDriver.SubscribeEventSources(cctx, conn, closeSubCh, depExpression, deps, filterFunc, actionFunc)
 							if err != nil {
@@ -184,6 +193,8 @@ func (sensorCtx *SensorContext) ListenEvents(ctx context.Context, stopCh <-chan
 	logger.Info("Sensor started.")
 	<-stopCh
 	logger.Info("Shutting down...")
+	cancel()
+	wg.Wait()
 	return nil
 }