From 57d1a099aaa26ea0da7ff70944b4f14566d75292 Mon Sep 17 00:00:00 2001
From: Peter Baumgartner <pete@lincolnloop.com>
Date: Mon, 25 Mar 2024 14:06:33 -0600
Subject: [PATCH] Allow `ps` on clusters with large task count

---
 CHANGELOG.md |  4 ++++
 app/app.go   | 37 +++++++++++++++++++++++++++----------
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7f9f34c..aafd714 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 * Upgraded to Go 1.22 and updated dependencies
 
+### Fixed
+
+* Allow `ps` to work on clusters with large active task counts
+
 ## [4.1.0] - 2023-01-10
 
 ### Added
diff --git a/app/app.go b/app/app.go
index 4d9a44d..4aa3f7d 100644
--- a/app/app.go
+++ b/app/app.go
@@ -27,6 +27,8 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+const maxEcsDescribeTaskCount = 100
+
 var (
 	maxLifetime    = 12 * 60 * 60
 	waitForConnect = 60
@@ -785,16 +787,26 @@ func (a *App) DescribeTasks() ([]*ecs.Task, error) {
 		return nil, err
 	}
 	ecsSvc := ecs.New(a.Session)
-	var taskARNs []*string
+	chunkedTaskARNs := [][]*string{{}}
 	input := ecs.ListTasksInput{
 		Cluster: &a.Settings.Cluster.ARN,
 	}
+	logrus.WithFields(logrus.Fields{"cluster": a.Settings.Cluster.ARN}).Debug("fetching task list")
+
+	// handle chunking logic
+	addTaskARNToChunk := func(taskARN *string) {
+		if len(chunkedTaskARNs[len(chunkedTaskARNs)-1]) >= maxEcsDescribeTaskCount {
+			chunkedTaskARNs = append(chunkedTaskARNs, []*string{})
+		}
+		chunkedTaskARNs[len(chunkedTaskARNs)-1] = append(chunkedTaskARNs[len(chunkedTaskARNs)-1], taskARN)
+	}
+
 	err = ecsSvc.ListTasksPages(&input, func(resp *ecs.ListTasksOutput, lastPage bool) bool {
 		for _, taskARN := range resp.TaskArns {
 			if taskARN == nil {
 				continue
 			}
-			taskARNs = append(taskARNs, taskARN)
+			addTaskARNToChunk(taskARN)
 		}
 
 		return !lastPage
@@ -802,16 +814,21 @@ func (a *App) DescribeTasks() ([]*ecs.Task, error) {
 	if err != nil {
 		return nil, err
 	}
-	describeTasksOutput, err := ecsSvc.DescribeTasks(&ecs.DescribeTasksInput{
-		Tasks:   taskARNs,
-		Cluster: &a.Settings.Cluster.ARN,
-		Include: []*string{aws.String("TAGS")},
-	})
-	if err != nil {
-		return nil, err
+	var describedTasks []*ecs.Task
+	for i := range chunkedTaskARNs {
+		logrus.WithFields(logrus.Fields{"count": len(chunkedTaskARNs[i])}).Debug("fetching task descriptions")
+		describeTasksOutput, err := ecsSvc.DescribeTasks(&ecs.DescribeTasksInput{
+			Tasks:   chunkedTaskARNs[i],
+			Cluster: &a.Settings.Cluster.ARN,
+			Include: []*string{aws.String("TAGS")},
+		})
+		if err != nil {
+			return nil, err
+		}
+		describedTasks = append(describedTasks, describeTasksOutput.Tasks...)
 	}
 	var appTasks []*ecs.Task
-	for _, task := range describeTasksOutput.Tasks {
+	for _, task := range describedTasks {
 		isApp := false
 		isReviewApp := false
 		for _, t := range task.Tags {