From c876462eb07a1bc2b3af23f819915bc5c2534082 Mon Sep 17 00:00:00 2001
From: Owen Marshall <dev@owenmarshall.com>
Date: Wed, 1 Mar 2017 13:06:34 -0500
Subject: [PATCH] Chunk taskArns into groups of 100

If the ECS cluster has > 100 tasks, passing them to
ecs.DescribeTasksRequest() will result in the AWS API returning
errors.

This patch breaks them into chunks of at most 100, and calls
DescribeTasks for each chunk.

We also return early in case ListTasks returns no values; this
prevents DescribeTasks from throwing HTTP errors.
---
 provider/ecs.go      | 47 ++++++++++++++++++++++++++++++++++++--------
 provider/ecs_test.go | 39 ++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 8 deletions(-)

diff --git a/provider/ecs.go b/provider/ecs.go
index a85c14c30..a8426b43d 100644
--- a/provider/ecs.go
+++ b/provider/ecs.go
@@ -206,13 +206,28 @@ func (provider *ECS) listInstances(ctx context.Context, client *awsClient) ([]ec
 		taskArns = append(taskArns, req.Data.(*ecs.ListTasksOutput).TaskArns...)
 	}
 
-	req, taskResp := client.ecs.DescribeTasksRequest(&ecs.DescribeTasksInput{
-		Tasks:   taskArns,
-		Cluster: &provider.Cluster,
-	})
+	// Early return: if we can't list tasks we have nothing to
+	// describe below - likely empty cluster/permissions are bad.  This
+	// stops the AWS API from returning a 401 when you DescribeTasks
+	// with no input.
+	if len(taskArns) == 0 {
+		return []ecsInstance{}, nil
+	}
+
+	chunkedTaskArns := provider.chunkedTaskArns(taskArns)
+	var tasks []*ecs.Task
+
+	for _, arns := range chunkedTaskArns {
+		req, taskResp := client.ecs.DescribeTasksRequest(&ecs.DescribeTasksInput{
+			Tasks:   arns,
+			Cluster: &provider.Cluster,
+		})
+
+		if err := wrapAws(ctx, req); err != nil {
+			return nil, err
+		}
+		tasks = append(tasks, taskResp.Tasks...)
 
-	if err := wrapAws(ctx, req); err != nil {
-		return nil, err
 	}
 
 	containerInstanceArns := make([]*string, 0)
@@ -221,7 +236,7 @@ func (provider *ECS) listInstances(ctx context.Context, client *awsClient) ([]ec
 	taskDefinitionArns := make([]*string, 0)
 	byTaskDefinition := make(map[string]int)
 
-	for _, task := range taskResp.Tasks {
+	for _, task := range tasks {
 		if _, found := byContainerInstance[*task.ContainerInstanceArn]; !found {
 			byContainerInstance[*task.ContainerInstanceArn] = len(containerInstanceArns)
 			containerInstanceArns = append(containerInstanceArns, task.ContainerInstanceArn)
@@ -243,7 +258,7 @@ func (provider *ECS) listInstances(ctx context.Context, client *awsClient) ([]ec
 	}
 
 	var instances []ecsInstance
-	for _, task := range taskResp.Tasks {
+	for _, task := range tasks {
 
 		machineIdx := byContainerInstance[*task.ContainerInstanceArn]
 		taskDefIdx := byTaskDefinition[*task.TaskDefinitionArn]
@@ -398,6 +413,22 @@ func (provider *ECS) getFrontendRule(i ecsInstance) string {
 	return "Host:" + strings.ToLower(strings.Replace(i.Name, "_", "-", -1)) + "." + provider.Domain
 }
 
+// ECS expects no more than 100 parameters be passed to a DescribeTask call; thus, pack
+// each string into an array capped at 100 elements
+func (provider *ECS) chunkedTaskArns(tasks []*string) [][]*string {
+	var chunkedTasks [][]*string
+	for i := 0; i < len(tasks); i += 100 {
+		sliceEnd := -1
+		if i+100 < len(tasks) {
+			sliceEnd = i + 100
+		} else {
+			sliceEnd = len(tasks)
+		}
+		chunkedTasks = append(chunkedTasks, tasks[i:sliceEnd])
+	}
+	return chunkedTasks
+}
+
 func (i ecsInstance) Protocol() string {
 	if label := i.label("traefik.protocol"); label != "" {
 		return label
diff --git a/provider/ecs_test.go b/provider/ecs_test.go
index 0365819b3..c7baaca1a 100644
--- a/provider/ecs_test.go
+++ b/provider/ecs_test.go
@@ -308,3 +308,42 @@ func TestFilterInstance(t *testing.T) {
 		}
 	}
 }
+
+func TestTaskChunking(t *testing.T) {
+	provider := &ECS{}
+
+	testval := "a"
+	cases := []struct {
+		count           int
+		expectedLengths []int
+	}{
+		{0, []int(nil)},
+		{1, []int{1}},
+		{99, []int{99}},
+		{100, []int{100}},
+		{101, []int{100, 1}},
+		{199, []int{100, 99}},
+		{200, []int{100, 100}},
+		{201, []int{100, 100, 1}},
+		{555, []int{100, 100, 100, 100, 100, 55}},
+		{1001, []int{100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 1}},
+	}
+
+	for _, c := range cases {
+		var tasks []*string
+		for v := 0; v < c.count; v++ {
+			tasks = append(tasks, &testval)
+		}
+
+		out := provider.chunkedTaskArns(tasks)
+		var outCount []int
+
+		for _, el := range out {
+			outCount = append(outCount, len(el))
+		}
+
+		if !reflect.DeepEqual(outCount, c.expectedLengths) {
+			t.Errorf("Chunking %d elements, expected %#v, got %#v", c.count, c.expectedLengths, outCount)
+		}
+	}
+}