From 685af493bf0b2a71b64827c7786bb07db900d124 Mon Sep 17 00:00:00 2001
From: Mike Lang <mikelang3000@gmail.com>
Date: Thu, 12 Jan 2017 11:37:23 -0800
Subject: [PATCH] ecs probe: Allow cache settings to be tweaked

---
 probe/awsecs/client.go   |  2 +-
 probe/awsecs/reporter.go | 12 ++++++++----
 prog/main.go             |  6 +++++-
 prog/probe.go            |  2 +-
 4 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/probe/awsecs/client.go b/probe/awsecs/client.go
index 6eeccee60..b21444013 100644
--- a/probe/awsecs/client.go
+++ b/probe/awsecs/client.go
@@ -60,7 +60,7 @@ type ecsInfo struct {
 	taskServiceMap map[string]string
 }
 
-func newClient(cluster string, cacheSize int, cacheExpiry time.Duration) (*ecsClient, error) {
+func newClient(cluster string, cacheSize int, cacheExpiry time.Duration) (ecsClient, error) {
 	sess := session.New()
 
 	region, err := ec2metadata.New(sess).Region()
diff --git a/probe/awsecs/reporter.go b/probe/awsecs/reporter.go
index 982f5d704..67d539989 100644
--- a/probe/awsecs/reporter.go
+++ b/probe/awsecs/reporter.go
@@ -78,13 +78,17 @@ func getLabelInfo(rpt report.Report) map[string]map[string]*taskLabelInfo {
 
 // Reporter implements Tagger, Reporter
 type Reporter struct {
-	clientsByCluster map[string]*ecsClient
+	clientsByCluster map[string]ecsClient
+	cacheSize        int
+	cacheExpiry      time.Duration
 }
 
 // Make creates a new Reporter
-func Make() Reporter {
+func Make(cacheSize int, cacheExpiry time.Duration) Reporter {
 	return Reporter{
-		clientsByCluster: map[string]*ecsClient{},
+		clientsByCluster: map[string]ecsClient{},
+		cacheSize:        cacheSize,
+		cacheExpiry:      cacheExpiry,
 	}
 }
 
@@ -101,7 +105,7 @@ func (r Reporter) Tag(rpt report.Report) (report.Report, error) {
 		if !ok {
 			log.Debugf("Creating new ECS client")
 			var err error
-			client, err = newClient(cluster, 1e6, time.Hour) // TODO remove these temporary magic values
+			client, err = newClient(cluster, r.cacheSize, r.cacheExpiry)
 			if err != nil {
 				return rpt, err
 			}
diff --git a/prog/main.go b/prog/main.go
index ac1ca13a7..ed4d65d28 100644
--- a/prog/main.go
+++ b/prog/main.go
@@ -104,7 +104,9 @@ type probeFlags struct {
 	kubernetesEnabled bool
 	kubernetesConfig  kubernetes.ClientConfig
 
-	ecsEnabled bool
+	ecsEnabled     bool
+	ecsCacheSize   int
+	ecsCacheExpiry time.Duration
 
 	weaveEnabled  bool
 	weaveAddr     string
@@ -287,6 +289,8 @@ func main() {
 
 	// AWS ECS
 	flag.BoolVar(&flags.probe.ecsEnabled, "probe.ecs", false, "Collect ecs-related attributes for containers on this node")
+	flag.IntVar(&flags.probe.ecsCacheSize, "probe.ecs.cache.size", 1024*1024, "Max size of cached info for each ECS cluster")
+	flag.DurationVar(&flags.probe.ecsCacheExpiry, "probe.ecs.cache.expiry", time.Hour, "How long to keep cached ECS info")
 
 	// Weave
 	flag.StringVar(&flags.probe.weaveAddr, "probe.weave.addr", "127.0.0.1:6784", "IP address & port of the Weave router")
diff --git a/prog/probe.go b/prog/probe.go
index fd183a30b..946e93889 100644
--- a/prog/probe.go
+++ b/prog/probe.go
@@ -206,7 +206,7 @@ func probeMain(flags probeFlags, targets []appclient.Target) {
 	}
 
 	if flags.ecsEnabled {
-		reporter := awsecs.Make()
+		reporter := awsecs.Make(flags.ecsCacheSize, flags.ecsCacheExpiry)
 		p.AddReporter(reporter)
 		p.AddTagger(reporter)
 	}