Fix TestSessionCache_watcher flakiness (#66442)

lib/web/sessions.go

@@ -639,6 +639,8 @@ type sessionCacheOptions struct {
 	proxySigner multiplexer.PROXYHeaderSigner
 	// See [sessionCache.sessionWatcherStartImmediately]. Used for testing.
 	sessionWatcherStartImmediately bool
+	// See [sessionCache.sessionWatcherInitializedChannel]. Used for testing.
+	sessionWatcherInitializedChannel chan struct{}
 	// See [sessionCache.sessionWatcherEventProcessedChannel]. Used for testing.
 	sessionWatcherEventProcessedChannel chan struct{}
 	logger                              *slog.Logger
@@ -663,19 +665,26 @@ func newSessionCache(ctx context.Context, config sessionCacheOptions) (*sessionC
 	}
 
 	cache := &sessionCache{
-		clusterName:                         clusterName.GetClusterName(),
-		proxyClient:                         config.proxyClient,
-		accessPoint:                         config.accessPoint,
-		sessions:                            make(map[string]*SessionContext),
-		resources:                           make(map[string]*sessionResources),
-		authServers:                         config.servers,
-		closer:                              utils.NewCloseBroadcaster(),
-		cipherSuites:                        config.cipherSuites,
-		log:                                 config.logger,
-		clock:                               config.clock,
-		sessionLingeringThreshold:           config.sessionLingeringThreshold,
-		proxySigner:                         config.proxySigner,
-		sessionWatcherStartImmediately:      config.sessionWatcherStartImmediately,
+		clusterName:                      clusterName.GetClusterName(),
+		proxyClient:                      config.proxyClient,
+		accessPoint:                      config.accessPoint,
+		sessions:                         make(map[string]*SessionContext),
+		resources:                        make(map[string]*sessionResources),
+		authServers:                      config.servers,
+		closer:                           utils.NewCloseBroadcaster(),
+		cipherSuites:                     config.cipherSuites,
+		log:                              config.logger,
+		clock:                            config.clock,
+		sessionLingeringThreshold:        config.sessionLingeringThreshold,
+		proxySigner:                      config.proxySigner,
+		sessionWatcherStartImmediately:   config.sessionWatcherStartImmediately,
+		sessionWatcherInitializedChannel: config.sessionWatcherInitializedChannel,
+		sessionWatcherMarkInitialized: sync.OnceFunc(func() {
+			c := config.sessionWatcherInitializedChannel
+			if c != nil {
+				close(c)
+			}
+		}),
 		sessionWatcherEventProcessedChannel: config.sessionWatcherEventProcessedChannel,
 		buildType:                           config.buildType,
 	}
@@ -729,7 +738,14 @@ type sessionCache struct {
 	// backoff used to start the WebSession watcher.
 	// Used for testing.
 	sessionWatcherStartImmediately bool
-
+	// sessionWatcherInitializedChannel is used to signal that the sessionWatcher
+	// received its first OpInit event and is ready to observe updates.
+	// May be nil.
+	// Used for testing.
+	sessionWatcherInitializedChannel chan struct{}
+	// sessionWatcherMarkInitialized safely closes
+	// sessionWatcherInitializedChannel.
+	sessionWatcherMarkInitialized func()
 	// sessionWatcherEventProcessedChannel is used to signal that the
 	// sessionWatcher processed an event.
 	// May be nil.
@@ -859,6 +875,10 @@ func (s *sessionCache) watchWebSessionsOnce(ctx context.Context, reset func()) e
 				"event", logutils.StringerAttr(event),
 			)
 
+			if event.Type == types.OpInit {
+				s.sessionWatcherMarkInitialized()
+				continue
+			}
 			if event.Type != types.OpPut {
 				continue // We only care about OpPut at the moment.
 			}

lib/web/sessions_test.go

@@ -196,6 +196,7 @@ func TestSessionCache_watcher(t *testing.T) {
 	// cancel is used to make sure the sessionCache stops cleanly.
 	ctx := t.Context()
 
+	initializedC := make(chan struct{})
 	processedC := make(chan struct{})
 	sessionCache, err := newSessionCache(ctx, sessionCacheOptions{
 		buildType:   testModules.BuildType(),
@@ -207,11 +208,20 @@ func TestSessionCache_watcher(t *testing.T) {
 		clock:                               clock,
 		sessionLingeringThreshold:           1 * time.Minute,
 		sessionWatcherStartImmediately:      true,
+		sessionWatcherInitializedChannel:    initializedC,
 		sessionWatcherEventProcessedChannel: processedC,
 	})
 	require.NoError(t, err, "newSessionCache() failed")
 	defer sessionCache.Close()
 
+	// Wait for watcher initialization. This guarantees that the watcher is ready
+	// to observe updates.
+	select {
+	case <-initializedC:
+	case <-time.After(10 * time.Second):
+		t.Fatal("Timed out waiting for sessionCache watcher initialization")
+	}
+
 	// Sanity check active sessions.
 	require.Zero(t,
 		sessionCache.ActiveSessions(),