Fix ECS rollback stage does not remove canary created tasks (#4465)

* Fix cannot remove canary created task Signed-off-by: ductnn <[email protected]> * Fix reuse clean function Signed-off-by: ductnn <[email protected]> * Update logs message for ECS routing stage rollback Signed-off-by: ductnn <[email protected]> --------- Signed-off-by: ductnn <[email protected]>
pipe-cd · Jul 6, 2023 · 1c2e5cb · 1c2e5cb
1 parent 04e6721
commit 1c2e5cb
Showing 1 changed file with 33 additions and 4 deletions.
diff --git a/pkg/app/piped/executor/ecs/rollback.go b/pkg/app/piped/executor/ecs/rollback.go
@@ -82,19 +82,19 @@ func (e *rollbackExecutor) ensureRollback(ctx context.Context) model.StageStatus
 		return model.StageStatus_STAGE_FAILURE
 	}
 
-	primary, _, ok := loadTargetGroups(&e.Input, appCfg, runningDS)
+	primary, canary, ok := loadTargetGroups(&e.Input, appCfg, runningDS)
 	if !ok {
 		return model.StageStatus_STAGE_FAILURE
 	}
 
-	if !rollback(ctx, &e.Input, platformProviderName, platformProviderCfg, taskDefinition, serviceDefinition, primary) {
+	if !rollback(ctx, &e.Input, platformProviderName, platformProviderCfg, taskDefinition, serviceDefinition, primary, canary) {
 		return model.StageStatus_STAGE_FAILURE
 	}
 
 	return model.StageStatus_STAGE_SUCCESS
 }
 
-func rollback(ctx context.Context, in *executor.Input, platformProviderName string, platformProviderCfg *config.PlatformProviderECSConfig, taskDefinition types.TaskDefinition, serviceDefinition types.Service, targetGroup *types.LoadBalancer) bool {
+func rollback(ctx context.Context, in *executor.Input, platformProviderName string, platformProviderCfg *config.PlatformProviderECSConfig, taskDefinition types.TaskDefinition, serviceDefinition types.Service, primaryTargetGroup *types.LoadBalancer, canaryTargetGroup *types.LoadBalancer) bool {
 	in.LogPersister.Infof("Start rollback the ECS service and task family: %s and %s to original stage", *serviceDefinition.ServiceName, *taskDefinition.Family)
 	client, err := provider.DefaultRegistry().Client(platformProviderName, platformProviderCfg, in.Logger)
 	if err != nil {
@@ -127,7 +127,7 @@ func rollback(ctx context.Context, in *executor.Input, platformProviderName stri
 	}
 
 	// On rolling back, the scale of desired tasks will be set to 100 (same as the original state).
-	taskSet, err := client.CreateTaskSet(ctx, *service, *td, targetGroup, 100)
+	taskSet, err := client.CreateTaskSet(ctx, *service, *td, primaryTargetGroup, 100)
 	if err != nil {
 		in.LogPersister.Errorf("Failed to create ECS task set %s: %v", *serviceDefinition.ServiceName, err)
 		return false
@@ -147,6 +147,35 @@ func rollback(ctx context.Context, in *executor.Input, platformProviderName stri
 		}
 	}
 
+	// Reset routing
+	routingTrafficCfg := provider.RoutingTrafficConfig{
+		{
+			TargetGroupArn: *primaryTargetGroup.TargetGroupArn,
+			Weight:         100,
+		},
+		{
+			TargetGroupArn: *canaryTargetGroup.TargetGroupArn,
+			Weight:         0,
+		},
+	}
+
+	currListenerArns, err := client.GetListenerArns(ctx, *primaryTargetGroup)
+	if err != nil {
+		in.LogPersister.Errorf("Failed to get current active listeners: %v", err)
+		return false
+	}
+
+	if err := client.ModifyListeners(ctx, currListenerArns, routingTrafficCfg); err != nil {
+		in.LogPersister.Errorf("Failed to routing traffic to PRIMARY variant: %v", err)
+		return false
+	}
+
+	// Delete Canary taskSet
+	if !clean(ctx, in, platformProviderName, platformProviderCfg) {
+		in.LogPersister.Error("Failed to delete CANARY TaskSet")
+		return false
+	}
+
 	in.LogPersister.Infof("Rolled back the ECS service %s and task definition %s configuration to original stage", *serviceDefinition.ServiceName, *taskDefinition.Family)
 	return true
 }