diff --git a/cmd/run.go b/cmd/run.go index a1a13558..54a165ad 100644 --- a/cmd/run.go +++ b/cmd/run.go @@ -170,6 +170,16 @@ This enables a CRD-driven approach to automated image updates with Argo CD. return err } + // isLeader is an atomic bool that will be set to true when the controller + // becomes the leader. + var isLeader atomic.Bool + + // Start a goroutine that sets isLeader to true when leadership is acquired. + go func() { + <-mgr.Elected() + isLeader.Store(true) + }() + // Add the CacheWarmer as a Runnable to the manager. setupLogger.Info("Adding cache warmer to the manager.") warmupState := &WarmupStatus{ @@ -231,11 +241,11 @@ This enables a CRD-driven approach to automated image updates with Argo CD. } if err := mgr.AddReadyzCheck("warmup-check", func(req *http.Request) error { - if !warmupState.isCacheWarmed.Load() { - // If the cache is not yet warmed, the check fails. + // If we are the leader, we are ready only if the cache is warmed. + if isLeader.Load() && !warmupState.isCacheWarmed.Load() { return fmt.Errorf("cache is not yet warmed") } - // Once warmed, the check passes. + // If we are not the leader, we are always ready. return nil }); err != nil { setupLogger.Error(err, "unable to set up ready check") @@ -420,6 +430,8 @@ func (ws *WebhookServerRunnable) Start(ctx context.Context) error { return ws.webhookServer.Start(ctx) } +// NeedLeaderElection tells the manager that this runnable should only be +// run on the leader replica. func (ws *WebhookServerRunnable) NeedLeaderElection() bool { return true } diff --git a/cmd/run_test.go b/cmd/run_test.go index f358eaa5..b1c74c34 100644 --- a/cmd/run_test.go +++ b/cmd/run_test.go @@ -7,6 +7,7 @@ import ( "net/http" "os" "strconv" + "sync/atomic" "testing" "time" @@ -340,3 +341,68 @@ func TestReadyzCheckWithWarmupStatus(t *testing.T) { assert.Equal(t, "cache is not yet warmed", err.Error()) }) } + +// Assisted-by: Gemini AI +// TestLeadershipAwareReadyzCheck verifies the behavior of the leadership-aware readiness probe. +func TestLeadershipAwareReadyzCheck(t *testing.T) { + testCases := []struct { + name string + isLeader bool + isWarmed bool + expectError bool + errorMsg string + }{ + { + name: "Leader and not warmed", + isLeader: true, + isWarmed: false, + expectError: true, + errorMsg: "cache is not yet warmed", + }, + { + name: "Leader and warmed", + isLeader: true, + isWarmed: true, + expectError: false, + }, + { + name: "Not leader and not warmed", + isLeader: false, + isWarmed: false, + expectError: false, + }, + { + name: "Not leader and warmed", + isLeader: false, + isWarmed: true, + expectError: false, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + var isLeader atomic.Bool + isLeader.Store(tc.isLeader) + + status := &WarmupStatus{Done: make(chan struct{})} + status.isCacheWarmed.Store(tc.isWarmed) + + // This is the readiness check logic from cmd/run.go + readinessCheck := func(req *http.Request) error { + if isLeader.Load() && !status.isCacheWarmed.Load() { + return fmt.Errorf("cache is not yet warmed") + } + return nil + } + + err := readinessCheck(nil) + + if tc.expectError { + assert.Error(t, err) + assert.Equal(t, tc.errorMsg, err.Error()) + } else { + assert.NoError(t, err) + } + }) + } +}