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

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 {
+		if err := controller.SetupWithManager(mgr); err != nil {
 			setupLog.Error(err, "unable to setup controllers")
-			os.Exit(1)
+			cancel()
+			return
 		}
 	}()
 

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"
+	"errors"
+	"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 !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected error to wrap context.Canceled, got: %v", err)
+	}
+}