From a2d020cef05e6045019bef6987314ba4a17c2120 Mon Sep 17 00:00:00 2001
From: Tim Curtis <timjohncurtis@gmail.com>
Date: Tue, 8 Feb 2022 11:21:48 +0000
Subject: [PATCH] Add cli param for retry counts accessing the aws API

---
 README.md |  3 ++-
 go.mod    |  1 +
 main.go   | 14 +++++++++++---
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 6f6d198..b162284 100644
--- a/README.md
+++ b/README.md
@@ -20,6 +20,8 @@ The command line parameters that can be used are:
   discovery information to (default "ecs_file_sd.yml")
 * -config.role-arn (string): ARN of the role to assume when scraping
   the AWS API (optional)
+* -config.aws-api-retries (int): how many times to retry accessing the AWS API
+  before giving up and throwing the error (Can help with Throttling Exceptions)
 * -config.server-name-label (string): Docker label to define the server name
   (default "PROMETHEUS_EXPORTER_SERVER_NAME")
 * -config.job-name-label (string): Docker label to define the job name
@@ -97,4 +99,3 @@ that every minute, and by default Prometheus will reload the
 file the minute it is written).  After reloading your Prometheus
 master configuration, this program will begin informing via
 the discovery file of new targets that Prometheus must scrape.
-
diff --git a/go.mod b/go.mod
index 448b80c..9de22d4 100644
--- a/go.mod
+++ b/go.mod
@@ -3,6 +3,7 @@ module github.com/teralytics/prometheus-ecs-discovery
 go 1.15
 
 require (
+	github.com/aws/aws-sdk-go-v2 v1.3.1
 	github.com/aws/aws-sdk-go-v2/config v1.1.4
 	github.com/aws/aws-sdk-go-v2/credentials v1.1.4
 	github.com/aws/aws-sdk-go-v2/service/ec2 v1.3.0
diff --git a/main.go b/main.go
index 9688232..2b8948b 100644
--- a/main.go
+++ b/main.go
@@ -25,6 +25,8 @@ import (
 	"strings"
 	"time"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/aws/retry"
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
@@ -57,10 +59,11 @@ var cluster = flag.String("config.cluster", "", "name of the cluster to scrape")
 var outFile = flag.String("config.write-to", "ecs_file_sd.yml", "path of file to write ECS service discovery information to")
 var interval = flag.Duration("config.scrape-interval", 60*time.Second, "interval at which to scrape the AWS API for ECS service discovery information")
 var times = flag.Int("config.scrape-times", 0, "how many times to scrape before exiting (0 = infinite)")
+var awsApiRetries = flag.Int("config.aws-api-retries", 5, "how many retries to attempt to contact AWS API before giving up")
 var roleArn = flag.String("config.role-arn", "", "ARN of the role to assume when scraping the AWS API (optional)")
 var prometheusPortLabel = flag.String("config.port-label", "PROMETHEUS_EXPORTER_PORT", "Docker label to define the scrape port of the application (if missing an application won't be scraped)")
 var prometheusPathLabel = flag.String("config.path-label", "PROMETHEUS_EXPORTER_PATH", "Docker label to define the scrape path of the application")
-var prometheusSchemeLabel= flag.String("config.scheme-label", "PROMETHEUS_EXPORTER_SCHEME", "Docker label to define the scheme of the target application")
+var prometheusSchemeLabel = flag.String("config.scheme-label", "PROMETHEUS_EXPORTER_SCHEME", "Docker label to define the scheme of the target application")
 var prometheusFilterLabel = flag.String("config.filter-label", "", "Docker label (and optionally value) to require to scrape the application")
 var prometheusServerNameLabel = flag.String("config.server-name-label", "PROMETHEUS_EXPORTER_SERVER_NAME", "Docker label to define the server name")
 var prometheusJobNameLabel = flag.String("config.job-name-label", "PROMETHEUS_EXPORTER_JOB_NAME", "Docker label to define the job name")
@@ -297,7 +300,7 @@ func (t *AugmentedTask) ExporterInformation() []*PrometheusTaskInfo {
 
 		scheme, ok = d.DockerLabels[*prometheusSchemeLabel]
 		if ok {
-		    labels.Scheme = scheme
+			labels.Scheme = scheme
 		}
 
 		ret = append(ret, &PrometheusTaskInfo{
@@ -604,7 +607,12 @@ func GetAugmentedTasks(svc *ecs.Client, svcec2 *ec2.Client, clusterArns []*strin
 func main() {
 	flag.Parse()
 
-	config, err := config.LoadDefaultConfig(context.Background())
+	config, err := config.LoadDefaultConfig(
+		context.Background(),
+		config.WithRetryer(func() aws.Retryer {
+			return retry.AddWithMaxAttempts(retry.NewStandard(), *awsApiRetries)
+		}),
+	)
 	if err != nil {
 		logError(err)
 		return