Support AWS LoadBalancer with multi listeners

Signed-off-by: khanhtc1202 <[email protected]>
pipe-cd · Jul 3, 2023 · cfcd3fb · cfcd3fb
1 parent 2621f44
commit cfcd3fb
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 31 deletions.
diff --git a/pkg/app/piped/executor/ecs/ecs.go b/pkg/app/piped/executor/ecs/ecs.go
@@ -409,13 +409,13 @@ func routing(ctx context.Context, in *executor.Input, platformProviderName strin
 		in.Logger.Error("Failed to store traffic routing config to metadata store", zap.Error(err))
 	}
 
-	currListenerArn, err := client.GetListener(ctx, primaryTargetGroup)
+	currListenersArn, err := client.GetListenersArn(ctx, primaryTargetGroup)
 	if err != nil {
 		in.LogPersister.Errorf("Failed to get current active listener: %v", err)
 		return false
 	}
 
-	if err := client.ModifyListener(ctx, currListenerArn, routingTrafficCfg); err != nil {
+	if err := client.ModifyListeners(ctx, currListenersArn, routingTrafficCfg); err != nil {
 		in.LogPersister.Errorf("Failed to routing traffic to CANARY variant: %v", err)
 		return false
 	}

diff --git a/pkg/app/piped/platformprovider/ecs/client.go b/pkg/app/piped/platformprovider/ecs/client.go
@@ -277,29 +277,29 @@ func (c *client) ServiceExists(ctx context.Context, clusterName string, serviceN
 	return false, nil
 }
 
-func (c *client) GetListener(ctx context.Context, targetGroup types.LoadBalancer) (string, error) {
+func (c *client) GetListenersArn(ctx context.Context, targetGroup types.LoadBalancer) ([]string, error) {
 	loadBalancerArn, err := c.getLoadBalancerArn(ctx, *targetGroup.TargetGroupArn)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 
 	input := &elasticloadbalancingv2.DescribeListenersInput{
 		LoadBalancerArn: aws.String(loadBalancerArn),
 	}
 	output, err := c.elbClient.DescribeListeners(ctx, input)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	if len(output.Listeners) == 0 {
-		return "", platformprovider.ErrNotFound
+		return nil, platformprovider.ErrNotFound
 	}
-	// Note: Suppose the load balancer only have one listener.
-	// TODO: Support multi listeners pattern.
-	if len(output.Listeners) > 1 {
-		return "", fmt.Errorf("invalid listener configuration pointed to %s target group", *targetGroup.TargetGroupArn)
+
+	arns := make([]string, 0, len(output.Listeners))
+	for i := range output.Listeners {
+		arns[i] = *output.Listeners[i].ListenerArn
 	}
 
-	return *output.Listeners[0].ListenerArn, nil
+	return arns, nil
 }
 
 func (c *client) getLoadBalancerArn(ctx context.Context, targetGroupArn string) (string, error) {
@@ -317,32 +317,42 @@ func (c *client) getLoadBalancerArn(ctx context.Context, targetGroupArn string)
 	return output.TargetGroups[0].LoadBalancerArns[0], nil
 }
 
-func (c *client) ModifyListener(ctx context.Context, listenerArn string, routingTrafficCfg RoutingTrafficConfig) error {
+func (c *client) ModifyListeners(ctx context.Context, listenersArn []string, routingTrafficCfg RoutingTrafficConfig) error {
 	if len(routingTrafficCfg) != 2 {
 		return fmt.Errorf("invalid listener configuration: requires 2 target groups")
 	}
-	input := &elasticloadbalancingv2.ModifyListenerInput{
-		ListenerArn: aws.String(listenerArn),
-		DefaultActions: []elbtypes.Action{
-			{
-				Type: elbtypes.ActionTypeEnumForward,
-				ForwardConfig: &elbtypes.ForwardActionConfig{
-					TargetGroups: []elbtypes.TargetGroupTuple{
-						{
-							TargetGroupArn: aws.String(routingTrafficCfg[0].TargetGroupArn),
-							Weight:         aws.Int32(int32(routingTrafficCfg[0].Weight)),
-						},
-						{
-							TargetGroupArn: aws.String(routingTrafficCfg[1].TargetGroupArn),
-							Weight:         aws.Int32(int32(routingTrafficCfg[1].Weight)),
+
+	modifyListener := func(ctx context.Context, listenerArn string) error {
+		input := &elasticloadbalancingv2.ModifyListenerInput{
+			ListenerArn: aws.String(listenerArn),
+			DefaultActions: []elbtypes.Action{
+				{
+					Type: elbtypes.ActionTypeEnumForward,
+					ForwardConfig: &elbtypes.ForwardActionConfig{
+						TargetGroups: []elbtypes.TargetGroupTuple{
+							{
+								TargetGroupArn: aws.String(routingTrafficCfg[0].TargetGroupArn),
+								Weight:         aws.Int32(int32(routingTrafficCfg[0].Weight)),
+							},
+							{
+								TargetGroupArn: aws.String(routingTrafficCfg[1].TargetGroupArn),
+								Weight:         aws.Int32(int32(routingTrafficCfg[1].Weight)),
+							},
 						},
 					},
 				},
 			},
-		},
+		}
+		_, err := c.elbClient.ModifyListener(ctx, input)
+		return err
+	}
+
+	for _, listener := range listenersArn {
+		if err := modifyListener(ctx, listener); err != nil {
+			return err
+		}
 	}
-	_, err := c.elbClient.ModifyListener(ctx, input)
-	return err
+	return nil
 }
 
 func (c *client) TagResource(ctx context.Context, resourceArn string, tags []types.Tag) error {

diff --git a/pkg/app/piped/platformprovider/ecs/ecs.go b/pkg/app/piped/platformprovider/ecs/ecs.go
@@ -55,8 +55,8 @@ type ECS interface {
 }
 
 type ELB interface {
-	GetListener(ctx context.Context, targetGroup types.LoadBalancer) (string, error)
-	ModifyListener(ctx context.Context, listenerArn string, routingTrafficCfg RoutingTrafficConfig) error
+	GetListenersArn(ctx context.Context, targetGroup types.LoadBalancer) ([]string, error)
+	ModifyListeners(ctx context.Context, listenersArn []string, routingTrafficCfg RoutingTrafficConfig) error
 }
 
 // Registry holds a pool of aws client wrappers.