Caching responses from EC2

Signed-off-by: Dave Henderson <dhenderson@gmail.com>

5 files changed, 84 insertions, 35 deletions

diff --git a/aws/ec2info.go b/aws/ec2info.go
index 649977a1..6a8d1b4a 100644
--- a/aws/ec2info.go
+++ b/aws/ec2info.go
@@ -10,6 +10,7 @@ import (
 type Ec2Info struct {
 	describer  func() InstanceDescriber
 	metaClient *Ec2Meta
+	cache      map[string]interface{}
 }
 
 // InstanceDescriber - A subset of ec2iface.EC2API that we can use to call EC2.DescribeInstances
@@ -26,6 +27,7 @@ func NewEc2Info() *Ec2Info {
 			return ec2Client(region)
 		},
 		metaClient: metaClient,
+		cache:      make(map[string]interface{}),
 	}
 }
 
@@ -38,19 +40,12 @@ func ec2Client(region string) (client InstanceDescriber) {
 
 // Tag -
 func (e *Ec2Info) Tag(tag string, def ...string) string {
-	if e.metaClient.nonAWS {
-		returnDefault(def)
-	}
-
-	instanceID := e.metaClient.Meta("instance-id")
-	input := &ec2.DescribeInstancesInput{
-		InstanceIds: aws.StringSlice([]string{instanceID}),
-	}
-	output, err := e.describer().DescribeInstances(input)
-	if err != nil {
+	output := e.describeInstance()
+	if output == nil {
 		return returnDefault(def)
 	}
-	if output != nil && len(output.Reservations) > 0 &&
+
+	if len(output.Reservations) > 0 &&
 		len(output.Reservations[0].Instances) > 0 &&
 		len(output.Reservations[0].Instances[0].Tags) > 0 {
 		for _, v := range output.Reservations[0].Instances[0].Tags {
@@ -62,3 +57,27 @@ func (e *Ec2Info) Tag(tag string, def ...string) string {
 
 	return returnDefault(def)
 }
+
+func (e *Ec2Info) describeInstance() (output *ec2.DescribeInstancesOutput) {
+	if e.metaClient.nonAWS {
+		return nil
+	}
+
+	if cached, ok := e.cache["DescribeInstances"]; ok {
+		output = cached.(*ec2.DescribeInstancesOutput)
+	} else {
+		instanceID := e.metaClient.Meta("instance-id")
+
+		input := &ec2.DescribeInstancesInput{
+			InstanceIds: aws.StringSlice([]string{instanceID}),
+		}
+
+		var err error
+		output, err = e.describer().DescribeInstances(input)
+		if err != nil {
+			return nil
+		}
+		e.cache["DescribeInstances"] = output
+	}
+	return
+}
diff --git a/aws/ec2info_test.go b/aws/ec2info_test.go
index 5911827d..1aac1907 100644
--- a/aws/ec2info_test.go
+++ b/aws/ec2info_test.go
@@ -8,6 +8,26 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+// test doubles
+type DummyInstanceDescriber struct {
+	tags []*ec2.Tag
+}
+
+func (d DummyInstanceDescriber) DescribeInstances(*ec2.DescribeInstancesInput) (*ec2.DescribeInstancesOutput, error) {
+	output := &ec2.DescribeInstancesOutput{
+		Reservations: []*ec2.Reservation{
+			&ec2.Reservation{
+				Instances: []*ec2.Instance{
+					&ec2.Instance{
+						Tags: d.tags,
+					},
+				},
+			},
+		},
+	}
+	return output, nil
+}
+
 func TestTag_MissingKey(t *testing.T) {
 	server, ec2meta := MockServer(200, `"i-1234"`)
 	defer server.Close()
@@ -28,6 +48,7 @@ func TestTag_MissingKey(t *testing.T) {
 			return client
 		},
 		metaClient: ec2meta,
+		cache:      make(map[string]interface{}),
 	}
 
 	assert.Empty(t, e.Tag("missing"))
@@ -54,6 +75,7 @@ func TestTag_ValidKey(t *testing.T) {
 			return client
 		},
 		metaClient: ec2meta,
+		cache:      make(map[string]interface{}),
 	}
 
 	assert.Equal(t, "bar", e.Tag("foo"))
@@ -62,6 +84,7 @@ func TestTag_ValidKey(t *testing.T) {
 
 func TestTag_NonEC2(t *testing.T) {
 	server, ec2meta := MockServer(404, "")
+	ec2meta.nonAWS = true
 	defer server.Close()
 	client := DummyInstanceDescriber{}
 	e := &Ec2Info{
@@ -69,28 +92,9 @@ func TestTag_NonEC2(t *testing.T) {
 			return client
 		},
 		metaClient: ec2meta,
+		cache:      make(map[string]interface{}),
 	}
 
 	assert.Equal(t, "", e.Tag("foo"))
 	assert.Equal(t, "default", e.Tag("foo", "default"))
 }
-
-// test doubles
-type DummyInstanceDescriber struct {
-	tags []*ec2.Tag
-}
-
-func (d DummyInstanceDescriber) DescribeInstances(*ec2.DescribeInstancesInput) (*ec2.DescribeInstancesOutput, error) {
-	output := &ec2.DescribeInstancesOutput{
-		Reservations: []*ec2.Reservation{
-			&ec2.Reservation{
-				Instances: []*ec2.Instance{
-					&ec2.Instance{
-						Tags: d.tags,
-					},
-				},
-			},
-		},
-	}
-	return output, nil
-}
diff --git a/aws/ec2meta.go b/aws/ec2meta.go
index b17112db..c441280e 100644
--- a/aws/ec2meta.go
+++ b/aws/ec2meta.go
@@ -17,6 +17,12 @@ type Ec2Meta struct {
 	Endpoint string
 	Client   *http.Client
 	nonAWS   bool
+	cache    map[string]string
+}
+
+// NewEc2Meta -
+func NewEc2Meta() *Ec2Meta {
+	return &Ec2Meta{cache: make(map[string]string)}
 }
 
 // returnDefault -
@@ -37,7 +43,11 @@ func unreachable(err error) bool {
 	return false
 }
 
-func (e *Ec2Meta) retrieveMetadata(url string, key string, def ...string) string {
+func (e *Ec2Meta) retrieveMetadata(url string, def ...string) string {
+	if value, ok := e.cache[url]; ok {
+		return value
+	}
+
 	if e.nonAWS {
 		return returnDefault(def)
 	}
@@ -63,6 +73,7 @@ func (e *Ec2Meta) retrieveMetadata(url string, key string, def ...string) string
 		log.Fatalf("Failed to read response body from %s: %v", url, err)
 	}
 	value := strings.TrimSpace(string(body))
+	e.cache[url] = value
 
 	return value
 }
@@ -74,7 +85,7 @@ func (e *Ec2Meta) Meta(key string, def ...string) string {
 	}
 
 	url := e.Endpoint + "/latest/meta-data/" + key
-	return e.retrieveMetadata(url, key, def...)
+	return e.retrieveMetadata(url, def...)
 }
 
 // Dynamic -
@@ -84,7 +95,7 @@ func (e *Ec2Meta) Dynamic(key string, def ...string) string {
 	}
 
 	url := e.Endpoint + "/latest/dynamic/" + key
-	return e.retrieveMetadata(url, key, def...)
+	return e.retrieveMetadata(url, def...)
 }
 
 // Region -
diff --git a/aws/ec2meta_test.go b/aws/ec2meta_test.go
index 0f1f04c8..f7a2e324 100644
--- a/aws/ec2meta_test.go
+++ b/aws/ec2meta_test.go
@@ -1,6 +1,7 @@
 package aws
 
 import (
+	"errors"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -58,3 +59,17 @@ func TestRegion_KnownRegion(t *testing.T) {
 
 	assert.Equal(t, "us-east-1", ec2meta.Region())
 }
+
+func TestUnreachable(t *testing.T) {
+	assert.False(t, unreachable(errors.New("foo")))
+	assert.True(t, unreachable(errors.New("host is down")))
+	assert.True(t, unreachable(errors.New("request canceled")))
+	assert.True(t, unreachable(errors.New("no route to host")))
+}
+
+func TestRetrieveMetadata_NonEC2(t *testing.T) {
+	ec2meta := NewEc2Meta()
+	ec2meta.nonAWS = true
+
+	assert.Equal(t, "foo", ec2meta.retrieveMetadata("", "foo"))
+}
diff --git a/aws/testutils.go b/aws/testutils.go
index 1a4d40b7..32c8d408 100644
--- a/aws/testutils.go
+++ b/aws/testutils.go
@@ -21,6 +21,6 @@ func MockServer(code int, body string) (*httptest.Server, *Ec2Meta) {
 	}
 	httpClient := &http.Client{Transport: tr}
 
-	client := &Ec2Meta{server.URL + "/", httpClient, false}
+	client := &Ec2Meta{server.URL + "/", httpClient, false, make(map[string]string)}
 	return server, client
 }