addressing review

Signed-off-by: Imran Pochi <[email protected]>

Author: ipochi

cmd/server/app/options/options.go

@@ -157,7 +157,7 @@ func (o *ProxyRunOptions) Flags() *pflag.FlagSet {
 	flags.BoolVar(&o.EnableLeaseController, "enable-lease-controller", o.EnableLeaseController, "Enable lease controller to publish and garbage collect proxy server leases.")
 	flags.StringVar(&o.LeaseNamespace, "lease-namespace", o.LeaseNamespace, "The namespace where lease objects are managed by the controller.")
 	flags.StringVar(&o.LeaseLabel, "lease-label", o.LeaseLabel, "The labels on which the lease objects are managed.")
-	flags.DurationVar(&o.GracefulShutdownTimeout, "graceful-shutdown-timeout", o.GracefulShutdownTimeout, "Timeout duration for graceful shutdown of the server. The server will wait for active connections to close before forcefully terminating.")
+	flags.DurationVar(&o.GracefulShutdownTimeout, "graceful-shutdown-timeout", o.GracefulShutdownTimeout, "Timeout duration for graceful shutdown of the server. The server will wait for active connections to close before forcefully terminating. Set to 0 to disable graceful shutdown (default: 0).")
 	flags.Bool("warn-on-channel-limit", true, "This behavior is now thread safe and always on. This flag will be removed in a future release.")
 	flags.MarkDeprecated("warn-on-channel-limit", "This behavior is now thread safe and always on. This flag will be removed in a future release.")
 
@@ -343,6 +343,11 @@ func (o *ProxyRunOptions) Validate() error {
 		}
 	}
 
+	// Validate graceful shutdown timeout
+	if o.GracefulShutdownTimeout < 0 {
+		return fmt.Errorf("graceful-shutdown-timeout must be >= 0, got %v", o.GracefulShutdownTimeout)
+	}
+
 	o.NeedsKubernetesClient = usingServiceAccountAuth || o.EnableLeaseController
 
 	return nil
@@ -386,7 +391,7 @@ func NewProxyRunOptions() *ProxyRunOptions {
 		EnableLeaseController:     false,
 		LeaseNamespace:            "kube-system",
 		LeaseLabel:                "k8s-app=konnectivity-server",
-		GracefulShutdownTimeout:   15 * time.Second,
+		GracefulShutdownTimeout:   0,
 	}
 	return &o
 }

cmd/server/app/options/options_test.go

@@ -63,6 +63,7 @@ func TestDefaultServerOptions(t *testing.T) {
 	assertDefaultValue(t, "CipherSuites", defaultServerOptions.CipherSuites, make([]string, 0))
 	assertDefaultValue(t, "XfrChannelSize", defaultServerOptions.XfrChannelSize, 10)
 	assertDefaultValue(t, "APIContentType", defaultServerOptions.APIContentType, "application/vnd.kubernetes.protobuf")
+	assertDefaultValue(t, "GracefulShutdownTimeout", defaultServerOptions.GracefulShutdownTimeout, 0*time.Second)
 
 }
 
@@ -168,6 +169,21 @@ func TestValidate(t *testing.T) {
 			value:    -10,
 			expected: fmt.Errorf("channel size -10 must be greater than 0"),
 		},
+		"NegativeGracefulShutdownTimeout": {
+			field:    "GracefulShutdownTimeout",
+			value:    -1 * time.Second,
+			expected: fmt.Errorf("graceful-shutdown-timeout must be >= 0, got -1s"),
+		},
+		"ZeroGracefulShutdownTimeout": {
+			field:    "GracefulShutdownTimeout",
+			value:    0 * time.Second,
+			expected: nil,
+		},
+		"PositiveGracefulShutdownTimeout": {
+			field:    "GracefulShutdownTimeout",
+			value:    30 * time.Second,
+			expected: nil,
+		},
 	} {
 		t.Run(desc, func(t *testing.T) {
 			testServerOptions := NewProxyRunOptions()
@@ -184,6 +200,10 @@ func TestValidate(t *testing.T) {
 				case reflect.Int:
 					ivalue := tc.value.(int)
 					fv.SetInt(int64(ivalue))
+				case reflect.Int64:
+					if duration, ok := tc.value.(time.Duration); ok {
+						fv.SetInt(int64(duration))
+					}
 				}
 			}
 			actual := testServerOptions.Validate()

cmd/server/app/server.go

@@ -186,7 +186,28 @@ func (p *Proxy) Run(o *options.ProxyRunOptions, stopCh <-chan struct{}) error {
 	}
 
 	<-stopCh
-	klog.V(1).Infoln("Received shutdown signal, initiating graceful shutdown.")
+	klog.V(1).Infoln("Shutting down server.")
+
+	// If graceful shutdown timeout is 0, use the old behavior (immediate shutdown)
+	if o.GracefulShutdownTimeout == 0 {
+		if frontendStop != nil {
+			if err := frontendStop(context.Background()); err != nil {
+				klog.ErrorS(err, "failed to stop frontend server")
+			}
+		}
+		if p.agentServer != nil {
+			p.agentServer.Stop()
+		}
+		if p.adminServer != nil {
+			p.adminServer.Close()
+		}
+		if p.healthServer != nil {
+			p.healthServer.Close()
+		}
+		return nil
+	}
+
+	klog.V(1).Infoln("Initiating graceful shutdown.")
 
 	// Start graceful shutdown with timeout
 	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), o.GracefulShutdownTimeout)
@@ -195,9 +216,21 @@ func (p *Proxy) Run(o *options.ProxyRunOptions, stopCh <-chan struct{}) error {
 	// Create a WaitGroup to track shutdown of all components
 	var wg sync.WaitGroup
 
-	// Shutdown frontend server gracefully (if available)
+	// Add all workers to WaitGroup upfront
 	if frontendStop != nil {
 		wg.Add(1)
+	}
+	wg.Add(3) // agent, admin, health servers
+
+	// Create completion channel before starting goroutines
+	shutdownComplete := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(shutdownComplete)
+	}()
+
+	// Shutdown frontend server gracefully (if available)
+	if frontendStop != nil {
 		go func() {
 			defer wg.Done()
 			klog.V(1).Infoln("Gracefully stopping frontend server...")
@@ -210,7 +243,6 @@ func (p *Proxy) Run(o *options.ProxyRunOptions, stopCh <-chan struct{}) error {
 	}
 
 	// Shutdown agent server gracefully
-	wg.Add(1)
 	go func() {
 		defer wg.Done()
 		klog.V(1).Infoln("Gracefully stopping agent server...")
@@ -219,20 +251,17 @@ func (p *Proxy) Run(o *options.ProxyRunOptions, stopCh <-chan struct{}) error {
 	}()
 
 	// Shutdown admin server gracefully
-	wg.Add(1)
 	go func() {
 		defer wg.Done()
 		klog.V(1).Infoln("Gracefully stopping admin server...")
 		if err := p.adminServer.Shutdown(shutdownCtx); err != nil {
 			klog.ErrorS(err, "failed to shut down admin server")
-
 		} else {
 			klog.V(1).Infoln("admin server stopped.")
 		}
 	}()
 
 	// Shutdown health server gracefully
-	wg.Add(1)
 	go func() {
 		defer wg.Done()
 		klog.V(1).Infoln("Gracefully stopping health server...")
@@ -244,12 +273,6 @@ func (p *Proxy) Run(o *options.ProxyRunOptions, stopCh <-chan struct{}) error {
 	}()
 
 	// Wait for all servers to shutdown or timeout
-	shutdownComplete := make(chan struct{})
-	go func() {
-		wg.Wait()
-		close(shutdownComplete)
-	}()
-
 	select {
 	case <-shutdownComplete:
 		klog.V(1).Infoln("Graceful shutdown completed successfully.")
@@ -268,8 +291,6 @@ func (p *Proxy) Run(o *options.ProxyRunOptions, stopCh <-chan struct{}) error {
 		// frontend server's force-stop is handled by its StopFunc
 	}
 
-	klog.V(1).Infoln("Server shutdown complete.")
-
 	return nil
 }