From f34c04ff993f9727ca8bdd73148a45a183c76343 Mon Sep 17 00:00:00 2001 From: Cam Hutchison Date: Thu, 25 Jun 2026 16:49:16 +1000 Subject: [PATCH] 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 | 12 ++++++------ lib/srv/discovery/eks_audit_log_watcher.go | 16 +++++++++++----- lib/srv/discovery/eks_audit_log_watcher_test.go | 11 +++++++++++ 3 files changed, 28 insertions(+), 11 deletions(-) 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) }) }