Merge pull request #12 from zmberg/fix-rbg-start-failed

Support for dynamic watch lws CRD

Author: gujingit

internal/controller/workloads/rolebasedgroup_controller.go

@@ -19,18 +19,17 @@ package workloads
 import (
 	"context"
 	"fmt"
-	apierrors "k8s.io/apimachinery/pkg/api/errors"
-	"k8s.io/apimachinery/pkg/util/errors"
 	"reflect"
-	lwsv1 "sigs.k8s.io/lws/api/leaderworkerset/v1"
-
-	"k8s.io/apimachinery/pkg/runtime/schema"
+	"sync"
 
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/errors"
 	"k8s.io/client-go/tools/record"
 	"k8s.io/klog/v2"
 	ctrl "sigs.k8s.io/controller-runtime"
@@ -40,12 +39,23 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/event"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/predicate"
+	"sigs.k8s.io/controller-runtime/pkg/reconcile"
+	lwsv1 "sigs.k8s.io/lws/api/leaderworkerset/v1"
 	workloadsv1alpha1 "sigs.k8s.io/rbgs/api/workloads/v1alpha1"
 	"sigs.k8s.io/rbgs/pkg/dependency"
 	"sigs.k8s.io/rbgs/pkg/reconciler"
 	"sigs.k8s.io/rbgs/pkg/utils"
 )
 
+var (
+	runtimeController *builder.TypedBuilder[reconcile.Request]
+	watchedWorkload   sync.Map
+)
+
+func init() {
+	watchedWorkload = sync.Map{}
+}
+
 // RoleBasedGroupReconciler reconciles a RoleBasedGroup object
 type RoleBasedGroupReconciler struct {
 	client    client.Client
@@ -94,6 +104,8 @@ func (r *RoleBasedGroupReconciler) Reconcile(ctx context.Context, req ctrl.Reque
 	var roleStatuses []workloadsv1alpha1.RoleStatus
 	var updateStatus bool
 	for _, role := range sortedRoles {
+		// first check whether watch lws cr
+		dynamicWatchCustomCRD(ctx, role.Workload)
 		// Check dependencies first
 		ready, err := dependencyManager.CheckDependencyReady(ctx, rbg, role)
 		if err != nil {
@@ -229,20 +241,21 @@ func (r *RoleBasedGroupReconciler) updateRBGStatus(ctx context.Context, rbg *wor
 
 // SetupWithManager sets up the controller with the Manager.
 func (r *RoleBasedGroupReconciler) SetupWithManager(mgr ctrl.Manager, options controller.Options) error {
-	controller := ctrl.NewControllerManagedBy(mgr).
+	runtimeController = ctrl.NewControllerManagedBy(mgr).
 		WithOptions(options).
 		For(&workloadsv1alpha1.RoleBasedGroup{}, builder.WithPredicates(RBGPredicate())).
 		Owns(&appsv1.StatefulSet{}, builder.WithPredicates(WorkloadPredicate())).
 		Owns(&appsv1.Deployment{}, builder.WithPredicates(WorkloadPredicate())).
 		Owns(&corev1.Service{}).
 		Named("workloads-rolebasedgroup")
 
-	err := utils.CheckCrdExists(r.apiReader, "leaderworkersets.leaderworkerset.x-k8s.io")
+	err := utils.CheckCrdExists(r.apiReader, reconciler.LwsCrdName)
 	if err == nil {
-		controller = controller.Owns(&lwsv1.LeaderWorkerSet{}, builder.WithPredicates(WorkloadPredicate()))
+		watchedWorkload.LoadOrStore(reconciler.LwsCrdName, struct{}{})
+		runtimeController.Owns(&lwsv1.LeaderWorkerSet{}, builder.WithPredicates(WorkloadPredicate()))
 	}
 
-	return controller.Complete(r)
+	return runtimeController.Complete(r)
 }
 
 // CheckCrdExists checks if the specified Custom Resource Definition (CRD) exists in the Kubernetes cluster.
@@ -357,3 +370,20 @@ func hasValidOwnerRef(obj client.Object, targetGVK schema.GroupVersionKind) bool
 func getRbgGVK() schema.GroupVersionKind {
 	return schema.FromAPIVersionAndKind(workloadsv1alpha1.GroupVersion.String(), "RoleBasedGroup")
 }
+
+func getLwsGVK() schema.GroupVersionKind {
+	return schema.FromAPIVersionAndKind(lwsv1.GroupVersion.String(), "LeaderWorkerSet")
+}
+
+func dynamicWatchCustomCRD(ctx context.Context, workload workloadsv1alpha1.WorkloadSpec) {
+	logger := log.FromContext(ctx)
+	switch workload.Kind {
+	case getLwsGVK().Kind:
+		_, lwsExist := watchedWorkload.Load(reconciler.LwsCrdName)
+		if !lwsExist {
+			watchedWorkload.LoadOrStore(reconciler.LwsCrdName, struct{}{})
+			runtimeController.Owns(&lwsv1.LeaderWorkerSet{}, builder.WithPredicates(WorkloadPredicate()))
+			logger.Info("rbgs controller watch LeaderWorkerSet CRD")
+		}
+	}
+}

pkg/reconciler/lws_reconciler.go

@@ -4,6 +4,9 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"reflect"
+	"time"
+
 	corev1 "k8s.io/api/core/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/runtime"
@@ -12,14 +15,16 @@ import (
 	"k8s.io/apimachinery/pkg/util/wait"
 	metaapplyv1 "k8s.io/client-go/applyconfigurations/meta/v1"
 	utilpointer "k8s.io/utils/pointer"
-	"reflect"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	lwsv1 "sigs.k8s.io/lws/api/leaderworkerset/v1"
 	lwsapplyv1 "sigs.k8s.io/lws/client-go/applyconfiguration/leaderworkerset/v1"
 	workloadsv1alpha1 "sigs.k8s.io/rbgs/api/workloads/v1alpha1"
 	"sigs.k8s.io/rbgs/pkg/utils"
-	"time"
+)
+
+const (
+	LwsCrdName = "leaderworkersets.leaderworkerset.x-k8s.io"
 )
 
 type LeaderWorkerSetReconciler struct {
@@ -103,6 +108,11 @@ func (r *LeaderWorkerSetReconciler) CheckWorkloadReady(ctx context.Context, rbg
 
 func (r *LeaderWorkerSetReconciler) CleanupOrphanedWorkloads(ctx context.Context, rbg *workloadsv1alpha1.RoleBasedGroup) error {
 	logger := log.FromContext(ctx)
+	err := utils.CheckCrdExists(r.client, LwsCrdName)
+	if err != nil {
+		logger.Info(fmt.Sprintf("LeaderWorkerSetReconciler CleanupOrphanedWorkloads check lws crd failed: %s", err.Error()))
+		return nil
+	}
 	// list lws managed by rbg
 	lwsList := &lwsv1.LeaderWorkerSetList{}
 	if err := r.client.List(ctx, lwsList, client.InNamespace(rbg.Namespace),