Adds ECS Cluster Region option

weaveworks · Sep 18, 2017 · cd21baf · cd21baf
1 parent b0cd180
commit cd21baf
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 11 deletions.
diff --git a/probe/awsecs/client.go b/probe/awsecs/client.go
@@ -70,16 +70,19 @@ 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, clusterRegion string) (EcsClient, error) {
 	sess := session.New()
+	var err error
 
-	region, err := ec2metadata.New(sess).Region()
-	if err != nil {
-		return nil, err
+	if clusterRegion == "" {
+		clusterRegion, err = ec2metadata.New(sess).Region()
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	return &ecsClientImpl{
-		client:       ecs.New(sess, &aws.Config{Region: aws.String(region)}),
+		client:       ecs.New(sess, &aws.Config{Region: aws.String(clusterRegion)}),
 		cluster:      cluster,
 		taskCache:    gcache.New(cacheSize).LRU().Expiration(cacheExpiry).Build(),
 		serviceCache: gcache.New(cacheSize).LRU().Expiration(cacheExpiry).Build(),

diff --git a/probe/awsecs/reporter.go b/probe/awsecs/reporter.go
@@ -87,16 +87,18 @@ type Reporter struct {
 	ClientsByCluster map[string]EcsClient // Exported for test
 	cacheSize        int
 	cacheExpiry      time.Duration
+	clusterRegion    string
 	handlerRegistry  *controls.HandlerRegistry
 	probeID          string
 }
 
 // Make creates a new Reporter
-func Make(cacheSize int, cacheExpiry time.Duration, handlerRegistry *controls.HandlerRegistry, probeID string) Reporter {
+func Make(cacheSize int, cacheExpiry time.Duration, clusterRegion string, handlerRegistry *controls.HandlerRegistry, probeID string) Reporter {
 	r := Reporter{
 		ClientsByCluster: map[string]EcsClient{},
 		cacheSize:        cacheSize,
 		cacheExpiry:      cacheExpiry,
+		clusterRegion:    clusterRegion,
 		handlerRegistry:  handlerRegistry,
 		probeID:          probeID,
 	}
@@ -114,7 +116,7 @@ func (r Reporter) getClient(cluster string) (EcsClient, error) {
 	if !ok {
 		log.Debugf("Creating new ECS client")
 		var err error
-		client, err = newClient(cluster, r.cacheSize, r.cacheExpiry)
+		client, err = newClient(cluster, r.cacheSize, r.cacheExpiry, r.clusterRegion)
 		if err != nil {
 			return nil, err
 		}

diff --git a/prog/main.go b/prog/main.go
@@ -124,9 +124,10 @@ type probeFlags struct {
 	kubernetesClientConfig kubernetes.ClientConfig
 	kubernetesKubeletPort  uint
 
-	ecsEnabled     bool
-	ecsCacheSize   int
-	ecsCacheExpiry time.Duration
+	ecsEnabled       bool
+	ecsCacheSize     int
+	ecsCacheExpiry   time.Duration
+	ecsClusterRegion string
 
 	weaveEnabled  bool
 	weaveAddr     string
@@ -323,6 +324,7 @@ func setupFlags(flags *flags) {
 	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")
+	flag.StringVar(&flags.probe.ecsClusterRegion, "probe.ecs.cluster.region", "", "ECS Cluster Region")
 
 	// 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
@@ -226,7 +226,7 @@ func probeMain(flags probeFlags, targets []appclient.Target) {
 	}
 
 	if flags.ecsEnabled {
-		reporter := awsecs.Make(flags.ecsCacheSize, flags.ecsCacheExpiry, handlerRegistry, probeID)
+		reporter := awsecs.Make(flags.ecsCacheSize, flags.ecsCacheExpiry, flags.ecsClusterRegion, handlerRegistry, probeID)
 		defer reporter.Stop()
 		p.AddReporter(reporter)
 		p.AddTagger(reporter)