From b8c4e1fca02b4587cb921cc655d6fda569a41841 Mon Sep 17 00:00:00 2001
From: Genevieve LEsperance <glesperance@pivotal.io>
Date: Fri, 4 May 2018 19:47:38 -0700
Subject: [PATCH] Add polling to instance deletion.

---
 aws/common/state.go      |  2 +-
 aws/ec2/instance.go      | 43 +++++++++++++++++++++++++++++++++++++++-
 aws/ec2/instance_test.go | 14 +++++++++++--
 aws/ec2/instances.go     |  2 +-
 4 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/aws/common/state.go b/aws/common/state.go
index 48bae241..02a4acce 100644
--- a/aws/common/state.go
+++ b/aws/common/state.go
@@ -36,7 +36,7 @@ func (s *State) Wait() (interface{}, error) {
 	notFoundChecks := 20
 	continuousTargetOccurence := 1
 	minTimeout := 2 * time.Second
-	delay := 10 * time.Second
+	delay := 2 * time.Second
 
 	type Result struct {
 		Result interface{}
diff --git a/aws/ec2/instance.go b/aws/ec2/instance.go
index 147d8fbc..4aa55f48 100644
--- a/aws/ec2/instance.go
+++ b/aws/ec2/instance.go
@@ -2,21 +2,25 @@ package ec2
 
 import (
 	"fmt"
+	"log"
 	"strings"
 
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/awserr"
 	awsec2 "github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/genevieve/leftovers/aws/common"
 )
 
 type Instance struct {
 	client       instancesClient
+	logger       logger
 	resourceTags resourceTags
 	id           *string
 	identifier   string
 	rtype        string
 }
 
-func NewInstance(client instancesClient, resourceTags resourceTags, id, keyName *string, tags []*awsec2.Tag) Instance {
+func NewInstance(client instancesClient, logger logger, resourceTags resourceTags, id, keyName *string, tags []*awsec2.Tag) Instance {
 	identifier := *id
 
 	extra := []string{}
@@ -34,6 +38,7 @@ func NewInstance(client instancesClient, resourceTags resourceTags, id, keyName
 
 	return Instance{
 		client:       client,
+		logger:       logger,
 		resourceTags: resourceTags,
 		id:           id,
 		identifier:   identifier,
@@ -41,6 +46,9 @@ func NewInstance(client instancesClient, resourceTags resourceTags, id, keyName
 	}
 }
 
+var pending = []string{"pending", "running", "shutting-down", "stopped", "stopping"}
+var target = []string{"terminated"}
+
 func (i Instance) Delete() error {
 	addresses, err := i.client.DescribeAddresses(&awsec2.DescribeAddressesInput{
 		Filters: []*awsec2.Filter{{
@@ -57,6 +65,14 @@ func (i Instance) Delete() error {
 		return fmt.Errorf("Terminate: %s", err)
 	}
 
+	refresh := instanceRefresh(i.client, i.id)
+	state := common.NewState(i.logger, refresh, pending, target)
+
+	_, err = state.Wait()
+	if err != nil {
+		return fmt.Errorf("Waiting for deletion: %s", err)
+	}
+
 	err = i.resourceTags.Delete("instance", *i.id)
 	if err != nil {
 		return fmt.Errorf("Delete resource tags: %s", err)
@@ -79,3 +95,28 @@ func (i Instance) Name() string {
 func (i Instance) Type() string {
 	return i.rtype
 }
+
+func instanceRefresh(client instancesClient, id *string) common.StateRefreshFunc {
+	return func() (interface{}, string, error) {
+		resp, err := client.DescribeInstances(&awsec2.DescribeInstancesInput{
+			InstanceIds: []*string{id},
+		})
+		if err != nil {
+			if ec2err, ok := err.(awserr.Error); ok && ec2err.Code() == "InvalidInstanceID.NotFound" {
+				resp = nil
+			} else {
+				log.Printf("Error on InstanceStateRefresh: %s", err)
+				return nil, "", err
+			}
+		}
+
+		if resp == nil || len(resp.Reservations) == 0 || len(resp.Reservations[0].Instances) == 0 {
+			return nil, "", nil
+		}
+
+		i := resp.Reservations[0].Instances[0]
+		state := *i.State.Name
+
+		return i, state, nil
+	}
+}
diff --git a/aws/ec2/instance_test.go b/aws/ec2/instance_test.go
index 2997cf24..3ec1969f 100644
--- a/aws/ec2/instance_test.go
+++ b/aws/ec2/instance_test.go
@@ -14,21 +14,24 @@ import (
 
 var _ = Describe("Instance", func() {
 	var (
-		instance     ec2.Instance
 		client       *fakes.InstancesClient
+		logger       *fakes.Logger
 		resourceTags *fakes.ResourceTags
 		id           *string
 		keyName      *string
+
+		instance ec2.Instance
 	)
 
 	BeforeEach(func() {
 		client = &fakes.InstancesClient{}
+		logger = &fakes.Logger{}
 		resourceTags = &fakes.ResourceTags{}
 		id = aws.String("the-id")
 		keyName = aws.String("the-key-name")
 		tags := []*awsec2.Tag{}
 
-		instance = ec2.NewInstance(client, resourceTags, id, keyName, tags)
+		instance = ec2.NewInstance(client, logger, resourceTags, id, keyName, tags)
 	})
 
 	Describe("Delete", func() {
@@ -38,6 +41,13 @@ var _ = Describe("Instance", func() {
 					AllocationId: aws.String("the-allocation-id"),
 				}},
 			}
+			client.DescribeInstancesCall.Returns.Output = &awsec2.DescribeInstancesOutput{
+				Reservations: []*awsec2.Reservation{{
+					Instances: []*awsec2.Instance{{
+						State: &awsec2.InstanceState{Name: aws.String("terminated")},
+					}},
+				}},
+			}
 		})
 
 		It("terminates the instance, deletes it's tags, and releases the address", func() {
diff --git a/aws/ec2/instances.go b/aws/ec2/instances.go
index d0c1a2bc..cb4131dd 100644
--- a/aws/ec2/instances.go
+++ b/aws/ec2/instances.go
@@ -45,7 +45,7 @@ func (i Instances) List(filter string) ([]common.Deletable, error) {
 	var resources []common.Deletable
 	for _, r := range instances.Reservations {
 		for _, instance := range r.Instances {
-			r := NewInstance(i.client, i.resourceTags, instance.InstanceId, instance.KeyName, instance.Tags)
+			r := NewInstance(i.client, i.logger, i.resourceTags, instance.InstanceId, instance.KeyName, instance.Tags)
 
 			if !strings.Contains(r.Name(), filter) {
 				continue