diff --git a/lib/srv/discovery/access_graph_aws.go b/lib/srv/discovery/access_graph_aws.go index e77d55c921b..c5774f89a74 100644 --- a/lib/srv/discovery/access_graph_aws.go +++ b/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) diff --git a/lib/srv/discovery/eks_audit_log_watcher.go b/lib/srv/discovery/eks_audit_log_watcher.go index 2688ddb6a92..ead35d4acb2 100644 --- a/lib/srv/discovery/eks_audit_log_watcher.go +++ b/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. } - }() + }) } } diff --git a/lib/srv/discovery/eks_audit_log_watcher_test.go b/lib/srv/discovery/eks_audit_log_watcher_test.go index 2815d8ce9d9..449b82b1fe9 100644 --- a/lib/srv/discovery/eks_audit_log_watcher_test.go +++ b/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) }) }