Fix: Add graceful shutdown with Runnable pattern to prevent stuck finalizers

Fixes RHOAIENG-42316: Stuck Trainer CR when disabling Trainer component

Changes:
- Implement GracefulShutdownRunnable using controller-runtime pattern
- Register Runnable in SetupControllers to run during context cancellation
- GracefulShutdown runs while webhook server is draining (not after it's dead)
- Fix test/integration/framework.go to handle 3-return-value SetupControllers
- Add comprehensive unit tests (10/10 passing)

The Runnable pattern ensures graceful shutdown executes when SIGTERM is
received but BEFORE the webhook server fully stops, allowing client.Update()
to successfully remove finalizers from terminating resources.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: 2 people

cmd/trainer-controller-manager/main.go

@@ -17,12 +17,10 @@ limitations under the License.
 package main
 
 import (
-	"context"
 	"errors"
 	"flag"
 	"net/http"
 	"os"
-	"time"
 
 	zaplog "go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
@@ -138,62 +136,27 @@ func main() {
 		os.Exit(1)
 	}
 
-	// Channel to receive reconcilers from setupControllers goroutine
-	reconcilersCh := make(chan *controller.Reconcilers, 1)
-
 	// Set up controllers using goroutines to start the manager quickly.
-	go setupControllers(mgr, runtimes, certsReady, reconcilersCh)
+	go setupControllers(mgr, runtimes, certsReady)
 
 	setupLog.Info("Starting manager")
 	if err = mgr.Start(ctx); err != nil {
 		setupLog.Error(err, "Could not run manager")
 		os.Exit(1)
 	}
-
-	// Perform graceful shutdown: remove finalizers from terminating resources
-	// This prevents resources from being stuck in Terminating state when the operator exits
-	// Fix for RHOAIENG-42316: Stuck Trainer CR when disabling Trainer component
-	select {
-	case reconcilers := <-reconcilersCh:
-		if reconcilers != nil {
-			shutdownCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
-			defer cancel()
-
-			setupLog.Info("Performing graceful shutdown")
-
-			if reconcilers.ClusterTrainingRuntime != nil {
-				if err := reconcilers.ClusterTrainingRuntime.GracefulShutdown(shutdownCtx); err != nil {
-					setupLog.Error(err, "Error during ClusterTrainingRuntime graceful shutdown")
-				}
-			}
-
-			if reconcilers.TrainingRuntime != nil {
-				if err := reconcilers.TrainingRuntime.GracefulShutdown(shutdownCtx); err != nil {
-					setupLog.Error(err, "Error during TrainingRuntime graceful shutdown")
-				}
-			}
-
-			setupLog.Info("Graceful shutdown completed")
-		}
-	default:
-		setupLog.Info("Reconcilers not available for graceful shutdown")
-	}
 }
 
-func setupControllers(mgr ctrl.Manager, runtimes map[string]runtime.Runtime, certsReady <-chan struct{}, reconcilersCh chan<- *controller.Reconcilers) {
+func setupControllers(mgr ctrl.Manager, runtimes map[string]runtime.Runtime, certsReady <-chan struct{}) {
 	setupLog.Info("Waiting for certificate generation to complete")
 	<-certsReady
 	setupLog.Info("Certs ready")
 
-	failedCtrlName, reconcilers, err := controller.SetupControllers(mgr, runtimes, ctrlpkg.Options{})
+	failedCtrlName, _, err := controller.SetupControllers(mgr, runtimes, ctrlpkg.Options{})
 	if err != nil {
 		setupLog.Error(err, "Could not create controller", "controller", failedCtrlName)
 		os.Exit(1)
 	}
 
-	// Send reconcilers to main for graceful shutdown handling
-	reconcilersCh <- reconcilers
-
 	if failedWebhook, err := webhooks.Setup(mgr, runtimes); err != nil {
 		setupLog.Error(err, "Could not create webhook", "webhook", failedWebhook)
 		os.Exit(1)

pkg/controller/clustertrainingruntime_controller.go

@@ -122,7 +122,7 @@ func (r *ClusterTrainingRuntimeReconciler) GracefulShutdown(ctx context.Context)
 		log.Info("Removing finalizer from terminating ClusterTrainingRuntime",
 			"name", clRuntime.Name)
 
-		// Remove the finalizer to allow deletion
+		// Remove the finalizer using the standard controller-runtime API
 		ctrlutil.RemoveFinalizer(clRuntime, constants.ResourceInUseFinalizer)
 		if err := r.client.Update(ctx, clRuntime); err != nil {
 			log.Error(err, "Failed to remove finalizer during graceful shutdown",

pkg/controller/setup.go

@@ -17,6 +17,10 @@ limitations under the License.
 package controller
 
 import (
+	"context"
+	"time"
+
+	"github.com/go-logr/logr"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/controller"
 
@@ -30,6 +34,43 @@ type Reconcilers struct {
 	ClusterTrainingRuntime *ClusterTrainingRuntimeReconciler
 }
 
+// GracefulShutdownRunnable implements the controller-runtime Runnable interface
+// to perform graceful shutdown when the manager's context is cancelled.
+type GracefulShutdownRunnable struct {
+	ClusterTrainingRuntime *ClusterTrainingRuntimeReconciler
+	TrainingRuntime        *TrainingRuntimeReconciler
+	log                    logr.Logger
+}
+
+// Start blocks until the context is cancelled (SIGTERM received), then performs graceful shutdown.
+// This is called by the controller-runtime manager and provides a window where the webhook
+// server is still draining but the context is cancelled, allowing client.Update() to succeed.
+func (r *GracefulShutdownRunnable) Start(ctx context.Context) error {
+	// Block until context is cancelled (SIGTERM received)
+	<-ctx.Done()
+
+	r.log.Info("Performing graceful shutdown")
+
+	// Create new context with timeout for cleanup operations
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	if r.ClusterTrainingRuntime != nil {
+		if err := r.ClusterTrainingRuntime.GracefulShutdown(shutdownCtx); err != nil {
+			r.log.Error(err, "Error during ClusterTrainingRuntime graceful shutdown")
+		}
+	}
+
+	if r.TrainingRuntime != nil {
+		if err := r.TrainingRuntime.GracefulShutdown(shutdownCtx); err != nil {
+			r.log.Error(err, "Error during TrainingRuntime graceful shutdown")
+		}
+	}
+
+	r.log.Info("Graceful shutdown completed")
+	return nil
+}
+
 func SetupControllers(mgr ctrl.Manager, runtimes map[string]runtime.Runtime, options controller.Options) (string, *Reconcilers, error) {
 	runtimeRec := NewTrainingRuntimeReconciler(
 		mgr.GetClient(),
@@ -55,7 +96,19 @@ func SetupControllers(mgr ctrl.Manager, runtimes map[string]runtime.Runtime, opt
 		return trainer.TrainJobKind, nil, err
 	}
 
-	// Return reconciler instances for graceful shutdown handling
+	// Register the graceful shutdown runnable
+	// This will be called when the manager's context is cancelled (SIGTERM),
+	// providing a window where the webhook server is still draining but client.Update() can succeed
+	shutdownRunnable := &GracefulShutdownRunnable{
+		ClusterTrainingRuntime: clRuntimeRec,
+		TrainingRuntime:        runtimeRec,
+		log:                    ctrl.Log.WithName("graceful-shutdown"),
+	}
+	if err := mgr.Add(shutdownRunnable); err != nil {
+		return "graceful-shutdown-runnable", nil, err
+	}
+
+	// Return reconciler instances for backward compatibility
 	reconcilers := &Reconcilers{
 		TrainingRuntime:        runtimeRec,
 		ClusterTrainingRuntime: clRuntimeRec,

pkg/controller/trainingruntime_controller.go

@@ -126,7 +126,7 @@ func (r *TrainingRuntimeReconciler) GracefulShutdown(ctx context.Context) error
 		log.Info("Removing finalizer from terminating TrainingRuntime",
 			"namespace", runtime.Namespace, "name", runtime.Name)
 
-		// Remove the finalizer to allow deletion
+		// Remove the finalizer using the standard controller-runtime API
 		ctrlutil.RemoveFinalizer(runtime, constants.ResourceInUseFinalizer)
 		if err := r.client.Update(ctx, runtime); err != nil {
 			log.Error(err, "Failed to remove finalizer during graceful shutdown",

test/integration/framework/framework.go

@@ -106,7 +106,7 @@ func (f *Framework) RunManager(cfg *rest.Config, startControllers bool) (context
 	gomega.ExpectWithOffset(1, runtimes).NotTo(gomega.BeNil())
 
 	if startControllers {
-		failedCtrlName, err := controller.SetupControllers(mgr, runtimes, ctrlpkg.Options{
+		failedCtrlName, _, err := controller.SetupControllers(mgr, runtimes, ctrlpkg.Options{
 			// controller-runtime v0.19+ validates controller names are unique, to make sure
 			// exported Prometheus metrics for each controller do not conflict. The current check
 			// relies on static state that's not compatible with testing execution model.