fix collaset revision panic: issue#148 (#154)

fix(CollaSet): fix collaset revision panic: issue #148
KusionStack · Feb 5, 2024 · eb06583 · eb06583
1 parent 4e14416
commit eb06583
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 9 deletions.
diff --git a/pkg/controllers/collaset/collaset_controller.go b/pkg/controllers/collaset/collaset_controller.go
@@ -70,7 +70,7 @@ func NewReconciler(mgr ctrl.Manager) reconcile.Reconciler {
 
 	return &CollaSetReconciler{
 		ReconcilerMixin: mixin,
-		revisionManager: revision.NewRevisionManager(mixin.Client, mixin.Scheme, &revisionOwnerAdapter{}),
+		revisionManager: revision.NewRevisionManager(mixin.Client, mixin.Scheme, NewRevisionOwnerAdapter(podcontrol.NewRealPodControl(mixin.Client, mixin.Scheme))),
 		syncControl:     synccontrol.NewRealSyncControl(mixin.Client, mixin.Logger, podcontrol.NewRealPodControl(mixin.Client, mixin.Scheme), mixin.Recorder),
 	}
 }

diff --git a/pkg/controllers/collaset/revision.go b/pkg/controllers/collaset/revision.go
@@ -19,9 +19,11 @@ package collaset
 import (
 	"encoding/json"
 
+	appsv1 "k8s.io/api/apps/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	appsv1alpha1 "kusionstack.io/operating/apis/apps/v1alpha1"
+	"kusionstack.io/operating/pkg/controllers/collaset/podcontrol"
 	"kusionstack.io/operating/pkg/controllers/utils/revision"
 )
 
@@ -53,9 +55,14 @@ func getCollaSetPatch(cls *appsv1alpha1.CollaSet) ([]byte, error) {
 	return patch, err
 }
 
-var _ revision.OwnerAdapter = &revisionOwnerAdapter{}
+func NewRevisionOwnerAdapter(podControl podcontrol.Interface) revision.OwnerAdapter {
+	return &revisionOwnerAdapter{
+		podControl: podControl,
+	}
+}
 
 type revisionOwnerAdapter struct {
+	podControl podcontrol.Interface
 }
 
 func (roa *revisionOwnerAdapter) GetSelector(obj metav1.Object) *metav1.LabelSelector {
@@ -83,6 +90,22 @@ func (roa *revisionOwnerAdapter) GetCurrentRevision(obj metav1.Object) string {
 	return ips.Status.CurrentRevision
 }
 
-func (roa *revisionOwnerAdapter) IsInUsed(_ metav1.Object, _ string) bool {
+func (roa *revisionOwnerAdapter) IsInUsed(obj metav1.Object, revision string) bool {
+	ips, _ := obj.(*appsv1alpha1.CollaSet)
+
+	if ips.Status.UpdatedRevision == revision || ips.Status.CurrentRevision == revision {
+		return true
+	}
+
+	pods, _ := roa.podControl.GetFilteredPods(ips.Spec.Selector, ips)
+	for _, pod := range pods {
+		if pod.Labels != nil {
+			currentRevisionName, exist := pod.Labels[appsv1.ControllerRevisionHashLabelKey]
+			if exist && currentRevisionName == revision {
+				return true
+			}
+		}
+	}
+
 	return false
 }
diff --git a/pkg/controllers/utils/revision/revision_manager.go b/pkg/controllers/utils/revision/revision_manager.go
@@ -194,12 +194,10 @@ func (rm *RevisionManager) cleanExpiredRevision(set metav1.Object, sortedRevisio
 		return sortedRevisions, nil
 	}
 
+	var cleanedRevisions []*apps.ControllerRevision
 	for _, revision := range *sortedRevisions {
-		if exceedNum == 0 {
-			break
-		}
-
-		if rm.ownerGetter.IsInUsed(set, revision.Name) {
+		if exceedNum == 0 || rm.ownerGetter.IsInUsed(set, revision.Name) {
+			cleanedRevisions = append(cleanedRevisions, revision)
 			continue
 		}
 
@@ -209,7 +207,6 @@ func (rm *RevisionManager) cleanExpiredRevision(set metav1.Object, sortedRevisio
 
 		exceedNum--
 	}
-	cleanedRevisions := (*sortedRevisions)[exceedNum:]
 
 	return &cleanedRevisions, nil
 }