fix: replace os.Exit in goroutine with context cancellation for graceful shutdown

The goroutine in main.go that waits for webhook readiness and sets up
controllers was calling os.Exit(1) on failure, which kills the process
immediately without releasing the leader election lease. Other replicas
cannot acquire leadership until the lease TTL expires (default 15s).
Additionally, WaitReady() was an infinite loop with no context or
timeout, risking a zombie leader state if the webhook health check
never passes.
This commit:
- Wraps the signal handler context with context.WithCancel
- Replaces os.Exit(1) in the goroutine with cancel() + return,
  triggering the manager's graceful shutdown path
- Updates WaitReady to accept a context.Context so it can exit
  when the context is cancelled
Fixes #2401

Signed-off-by: abhaygoudannavar <abhaysgoudnvr@gmail.com>

Author: abhaygoudannavar

main.go

@@ -17,6 +17,7 @@ limitations under the License.
 package main
 
 import (
+	"context"
 	"flag"
 	"net/http"
 	_ "net/http/pprof"
@@ -164,7 +165,9 @@ func main() {
 		})
 	}
 
-	ctx := ctrl.SetupSignalHandler()
+	signalCtx := ctrl.SetupSignalHandler()
+	ctx, cancel := context.WithCancel(signalCtx)
+	defer cancel()
 	cfg := ctrl.GetConfigOrDie()
 	setRestConfig(cfg)
 	cfg.UserAgent = "kruise-manager"
@@ -256,15 +259,17 @@ func main() {
 
 	go func() {
 		setupLog.Info("wait webhook ready")
-		if err = webhook.WaitReady(); err != nil {
+		if err = webhook.WaitReady(ctx); err != nil {
 			setupLog.Error(err, "unable to wait webhook ready")
-			os.Exit(1)
+			cancel()
+			return
 		}
 
 		setupLog.Info("setup controllers")
 		if err = controller.SetupWithManager(mgr); err != nil {
 			setupLog.Error(err, "unable to setup controllers")
-			os.Exit(1)
+			cancel()
+			return
 		}
 	}()
 

pkg/util/inplaceupdate/inplace_update_test.go

@@ -648,3 +648,65 @@ func TestRefresh(t *testing.T) {
 		}
 	}
 }
+
+func TestUpdateNextBatchReadsFromClone(t *testing.T) {
+	aHourAgo := metav1.NewTime(time.Unix(time.Now().Add(-time.Hour).Unix(), 0))
+	
+	oldClock := Clock
+	Clock = testingclock.NewFakeClock(aHourAgo.Time)
+	defer func() {
+		Clock = oldClock
+	}()
+
+	// storedPod is the pod in the fake client (what GetPod returns as "clone").
+	// It does NOT have the InPlaceUpdateState annotation.
+	storedPod := &v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name: "pod-stale-test",
+			Namespace: "default", // Fixed: Added namespace
+			Labels: map[string]string{
+				apps.StatefulSetRevisionLabel: "new-revision",
+			},
+		},
+		Spec: v1.PodSpec{
+			Containers:     []v1.Container{{Name: "c1", Image: "c1-img2"}, {Name: "c2", Image: "c2-img1"}},
+			ReadinessGates: []v1.PodReadinessGate{{ConditionType: appspub.InPlaceUpdateReady}},
+		},
+		Status: v1.PodStatus{
+			ContainerStatuses: []v1.ContainerStatus{
+				{Name: "c1", ImageID: "c1-img2-ID"},
+				{Name: "c2", ImageID: "c2-img1-ID"},
+			},
+		},
+	}
+
+	// stalePod is what gets passed to Refresh — it's a stale snapshot that still
+	// has the InPlaceUpdateState annotation with NextContainerImages.
+	stalePod := storedPod.DeepCopy()
+	stalePod.Annotations = map[string]string{
+		appspub.InPlaceUpdateStateKey: util.DumpJSON(appspub.InPlaceUpdateState{
+			Revision:            "new-revision",
+			UpdateTimestamp:     aHourAgo,
+			NextContainerImages: map[string]string{"c2": "c2-img2"},
+		}),
+	}
+
+	cli := fake.NewClientBuilder().WithObjects(storedPod).Build()
+	ctrl := New(cli, revisionadapter.NewDefaultImpl())
+
+	res := ctrl.Refresh(stalePod, nil)
+	if res.RefreshErr != nil {
+		t.Fatalf("Refresh should not fail, got: %v", res.RefreshErr)
+	}
+
+	// Verify the stored pod was NOT modified — updateNextBatch should have
+	// returned early because the clone has no InPlaceUpdateState annotation.
+	got := &v1.Pod{}
+	// Fixed: Include Namespace in types.NamespacedName
+	if err := cli.Get(context.TODO(), types.NamespacedName{Namespace: storedPod.Namespace, Name: storedPod.Name}, got); err != nil {
+		t.Fatalf("failed to get pod: %v", err)
+	}
+	if got.Spec.Containers[1].Image != "c2-img1" {
+		t.Fatalf("expected c2 image to remain c2-img1 (unchanged), got %s", got.Spec.Containers[1].Image)
+	}
+}

pkg/webhook/server.go

@@ -133,10 +133,16 @@ func Checker(req *http.Request) error {
 	return health.Checker(req)
 }
 
-func WaitReady() error {
+func WaitReady(ctx context.Context) error {
 	startTS := time.Now()
 	var err error
 	for {
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("context canceled while waiting for webhook ready: %w", ctx.Err())
+		default:
+		}
+
 		duration := time.Since(startTS)
 		if err = Checker(nil); err == nil {
 			return nil
@@ -147,5 +153,4 @@ func WaitReady() error {
 		}
 		time.Sleep(time.Second * 2)
 	}
-
 }

pkg/webhook/server_test.go

@@ -0,0 +1,36 @@
+/*
+Copyright 2024 The Kruise Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package webhook
+
+import (
+	"context"
+	"strings"
+	"testing"
+)
+
+func TestWaitReadyCancel(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // instantly cancel the context
+
+	err := WaitReady(ctx)
+	if err == nil {
+		t.Fatalf("expected error, got nil")
+	}
+	if !strings.Contains(err.Error(), "context canceled while waiting for webhook ready") {
+		t.Fatalf("unexpected error message: %v", err)
+	}
+}