supoort patch batch id in cacary release (#251)

Signed-off-by: zhihao jian <zhihao.jian@shopee.com>
Co-authored-by: zhihao jian <zhihao.jian@shopee.com>

pkg/controller/batchrelease/control/canarystyle/control_plane.go

@@ -103,11 +103,19 @@ func (rc *realCanaryController) UpgradeBatch() error {
 		return fmt.Errorf("wait canary workload %v reconcile", canary.GetCanaryInfo().LogKey)
 	}
 
-	batchContext := rc.CalculateBatchContext(rc.release)
+	batchContext, err := rc.CalculateBatchContext(rc.release)
+	if err != nil {
+		return err
+	}
 	klog.Infof("BatchRelease %v calculated context when upgrade batch: %s",
 		klog.KObj(rc.release), batchContext.Log())
 
-	return canary.UpgradeBatch(batchContext)
+	err = canary.UpgradeBatch(batchContext)
+	if err != nil {
+		return err
+	}
+
+	return rc.patcher.PatchPodBatchLabel(batchContext)
 }
 
 func (rc *realCanaryController) CheckBatchReady() error {
@@ -129,7 +137,10 @@ func (rc *realCanaryController) CheckBatchReady() error {
 		return fmt.Errorf("wait canary workload %v reconcile", canary.GetCanaryInfo().LogKey)
 	}
 
-	batchContext := rc.CalculateBatchContext(rc.release)
+	batchContext, err := rc.CalculateBatchContext(rc.release)
+	if err != nil {
+		return err
+	}
 	klog.Infof("BatchRelease %v calculated context when check batch ready: %s",
 		klog.KObj(rc.release), batchContext.Log())
 

pkg/controller/batchrelease/control/canarystyle/deployment/canary.go

@@ -43,6 +43,7 @@ type realCanaryController struct {
 	canaryObject *apps.Deployment
 	canaryClient client.Client
 	objectKey    types.NamespacedName
+	canaryPods   []*corev1.Pod
 }
 
 func newCanary(cli client.Client, key types.NamespacedName) realCanaryController {

pkg/controller/batchrelease/control/canarystyle/deployment/control.go

@@ -27,6 +27,7 @@ import (
 	"github.com/openkruise/rollouts/pkg/util"
 	utilclient "github.com/openkruise/rollouts/pkg/util/client"
 	apps "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -82,19 +83,30 @@ func (rc *realController) BuildCanaryController(release *v1beta1.BatchRelease) (
 	return rc, nil
 }
 
-func (rc *realController) CalculateBatchContext(release *v1beta1.BatchRelease) *batchcontext.BatchContext {
+func (rc *realController) CalculateBatchContext(release *v1beta1.BatchRelease) (*batchcontext.BatchContext, error) {
+	rolloutID := release.Spec.ReleasePlan.RolloutID
+	if rolloutID != "" {
+		// if rollout-id is set, the pod will be patched batch label,
+		// so we have to list pod here.
+		if _, err := rc.ListOwnedPods(); err != nil {
+			return nil, err
+		}
+	}
 	replicas := *rc.stableObject.Spec.Replicas
 	currentBatch := release.Status.CanaryStatus.CurrentBatch
 	desiredUpdate := int32(control.CalculateBatchReplicas(release, int(replicas), int(currentBatch)))
 
 	return &batchcontext.BatchContext{
+		Pods:                   rc.canaryPods,
+		RolloutID:              rolloutID,
 		Replicas:               replicas,
+		UpdateRevision:         release.Status.UpdateRevision,
 		CurrentBatch:           currentBatch,
 		DesiredUpdatedReplicas: desiredUpdate,
 		FailureThreshold:       release.Spec.ReleasePlan.FailureThreshold,
 		UpdatedReplicas:        rc.canaryObject.Status.Replicas,
 		UpdatedReadyReplicas:   rc.canaryObject.Status.AvailableReplicas,
-	}
+	}, nil
 }
 
 func (rc *realController) getLatestTemplate() (*v1.PodTemplateSpec, error) {
@@ -104,3 +116,12 @@ func (rc *realController) getLatestTemplate() (*v1.PodTemplateSpec, error) {
 	}
 	return &rc.stableObject.Spec.Template, nil
 }
+
+func (rc *realController) ListOwnedPods() ([]*corev1.Pod, error) {
+	if rc.canaryPods != nil {
+		return rc.canaryPods, nil
+	}
+	var err error
+	rc.canaryPods, err = util.ListOwnedPods(rc.canaryClient, rc.canaryObject)
+	return rc.canaryPods, err
+}

pkg/controller/batchrelease/control/canarystyle/deployment/control_test.go

@@ -226,7 +226,9 @@ func TestCalculateBatchContext(t *testing.T) {
 					canaryObject: canary,
 				},
 			}
-			got := control.CalculateBatchContext(cs.release())
+			got, err := control.CalculateBatchContext(cs.release())
+			got.FilterFunc = nil
+			Expect(err).NotTo(HaveOccurred())
 			Expect(reflect.DeepEqual(got, cs.result)).Should(BeTrue())
 		})
 	}
@@ -290,7 +292,8 @@ func TestRealCanaryController(t *testing.T) {
 	Expect(util.EqualIgnoreHash(&c.canaryObject.Spec.Template, &deployment.Spec.Template)).Should(BeTrue())
 
 	// check rolling
-	batchContext := c.CalculateBatchContext(release)
+	batchContext, err := c.CalculateBatchContext(release)
+	Expect(err).NotTo(HaveOccurred())
 	err = controller.UpgradeBatch(batchContext)
 	Expect(err).NotTo(HaveOccurred())
 	canary := getCanaryDeployment(release, deployment, c)

pkg/controller/batchrelease/control/canarystyle/interface.go

@@ -33,7 +33,7 @@ type Interface interface {
 	BuildCanaryController(release *v1beta1.BatchRelease) (CanaryInterface, error)
 	// CalculateBatchContext calculate the current batch context according to
 	// our release plan and the statues of stable workload and canary workload.
-	CalculateBatchContext(release *v1beta1.BatchRelease) *batchcontext.BatchContext
+	CalculateBatchContext(release *v1beta1.BatchRelease) (*batchcontext.BatchContext, error)
 }
 
 // CanaryInterface contains the methods about canary workload