discovery: Tie lifetime of EKS audit log fetcher to main function (#67962)

Ensure that the EKSAuditLogWatcher fully completes before
initializeAndWatchAccessGraph() completes, including the fetchers
spawned by the watcher.

This is done by adding the run of the EKSAuditLogWatcher to the
waitgroup the initialize function uses to wait for the grpc connection
to complete. The function does not return until this waitgroup is
finished.

Also add a WaitGroup to the EKSAuditLogWatcher to keep track of all the
fetchers spawned so that the EKSAuditLogWatcher.Run() function does not
return until all the fetchers have completed. This stops the fetchers
from trying to use a closed grpc connection during cleanup.

lib/srv/discovery/access_graph_aws.go

@@ -432,14 +432,12 @@ func (s *Server) initializeAndWatchAccessGraph(ctx context.Context, reloadCh <-c
 	// Start a goroutine to watch the access graph service connection state.
 	// If the connection is closed, cancel the context to stop the event watcher
 	// before it tries to send any events to the access graph service.
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
+	wg.Go(func() {
 		defer cancel()
 		if !accessGraphConn.WaitForStateChange(ctx, connectivity.Ready) {
 			s.Log.InfoContext(ctx, "Access graph service connection was closed")
 		}
-	}()
+	})
 
 	// Configure the poll interval
 	tickerInterval := defaultPollInterval
@@ -456,9 +454,11 @@ func (s *Server) initializeAndWatchAccessGraph(ctx context.Context, reloadCh <-c
 	s.Log.InfoContext(ctx, "Access graph service poll interval", "poll_interval", tickerInterval)
 
 	// Start the EKS audit log watcher that keeps track of the EKS audit log
-	// fetchers and updates them when Reconcile is called.
+	// fetchers and updates them when Reconcile is called. Use the WaitGroup
+	// to ensure this function does not return until all the log fetchers
+	// spawned from the watcher have completed.
 	eksAuditLogWatcher := newEKSAuditLogWatcher(client, s.Log)
-	go eksAuditLogWatcher.Run(ctx)
+	wg.Go(func() { eksAuditLogWatcher.Run(ctx) })
 
 	currentTAGResources := &aws_sync.Resources{}
 	timer := time.NewTimer(tickerInterval)

lib/srv/discovery/eks_audit_log_watcher.go

@@ -23,6 +23,7 @@ import (
 	"errors"
 	"iter"
 	"log/slog"
+	"sync"
 	"time"
 
 	"github.com/gravitational/trace"
@@ -122,6 +123,10 @@ func (w *eksAuditLogWatcher) Run(ctx context.Context) error {
 	w.log.InfoContext(ctx, "EKS Audit Log Watcher started")
 	defer w.log.InfoContext(ctx, "EKS Audit Log Watcher completed")
 
+	// Wait for all the fetchers we started to complete before we return.
+	var wg sync.WaitGroup
+	defer wg.Wait()
+
 	// initialize the stream but ignore errors. We do not want to exit on error
 	// as we need to ensure we keep the w.auditLogCLustersCh channel drained
 	// so as to not block AWS discovery.
@@ -138,7 +143,7 @@ func (w *eksAuditLogWatcher) Run(ctx context.Context) error {
 		select {
 		case clusters := <-w.auditLogClustersCh:
 			if stream != nil {
-				w.reconcile(ctx, clusters, stream)
+				w.reconcile(ctx, wg.Go, clusters, stream)
 			}
 		case completed := <-w.completedCh:
 			w.complete(ctx, completed)
@@ -202,11 +207,12 @@ func (w *eksAuditLogWatcher) Reconcile(ctx context.Context, clusters []eksAuditL
 // reconcile compares the given slice of clusters against the currently running
 // log fetchers and stops any running fetchers not in the cluster slice and
 // starts a log fetcher for any cluster in the slice that does not have a
-// running log fetcher.
+// running log fetcher. The spawn function is used to launch the goroutine that
+// manages the fetcher, typically WaitGroup.Go from the caller.
 //
 // Log fetchers that are started are initialized with the given grpc stream
 // over which they should send their audit logs.
-func (w *eksAuditLogWatcher) reconcile(ctx context.Context, clusters []eksAuditLogCluster, stream accessgraphv1alpha.AccessGraphService_KubeAuditLogStreamClient) {
+func (w *eksAuditLogWatcher) reconcile(ctx context.Context, spawn func(func()), clusters []eksAuditLogCluster, stream accessgraphv1alpha.AccessGraphService_KubeAuditLogStreamClient) {
 	w.log.DebugContext(ctx, "Reconciling EKS audit log clusters", "new_count", len(clusters))
 
 	// Make a map of the discovered clusters, keyed by ARN so we can compare against
@@ -233,7 +239,7 @@ func (w *eksAuditLogWatcher) reconcile(ctx context.Context, clusters []eksAuditL
 			logFetcher = w.newFetcher(discovered.fetcher, discovered.cluster, stream, w.log)
 		}
 		w.fetchers[arn] = cancel
-		go func() {
+		spawn(func() {
 			err := logFetcher.Run(fetcherCtx)
 			select {
 			case w.completedCh <- fetcherCompleted{arn, err}:
@@ -245,7 +251,7 @@ func (w *eksAuditLogWatcher) reconcile(ctx context.Context, clusters []eksAuditL
 				// specific context, which comes back to us via the case above
 				// on the completedCh channel, not here.
 			}
-		}()
+		})
 	}
 }
 

lib/srv/discovery/eks_audit_log_watcher_test.go

@@ -135,9 +135,20 @@ func TestEKSAuditLogWatcher_Reconcile(t *testing.T) {
 		require.Equal(t, 3, fetcherTracker.newCount)
 		require.True(t, f1.done)
 
+		// Add the two clusters back
+		watcher.Reconcile(ctx, []eksAuditLogCluster{{fetcher1, cluster1}, {fetcher2, cluster2}})
+		synctest.Wait()
+		require.Len(t, fetcherTracker.fetchers, 2)
+		require.Equal(t, 5, fetcherTracker.newCount)
+		require.True(t, f1.done)
+		require.True(t, f2.done)
+
+		// Cancel the watcher and ensure it waits for the fetchers to be done
 		cancel()
 		synctest.Wait()
 		require.ErrorIs(t, err, context.Canceled)
+		require.True(t, f1.done)
+		require.True(t, f2.done)
 	})
 }