From f446c1a1b73b85893ffe2cab96ab6661c054162a Mon Sep 17 00:00:00 2001
From: Rickard von Essen <rickard.von.essen@gmail.com>
Date: Thu, 8 Sep 2016 10:32:43 +0200
Subject: [PATCH] builders/digitalocean: fixes timeout waiting for snapshot
 #3853

---
 builder/digitalocean/step_snapshot.go | 23 ++++++------
 builder/digitalocean/wait.go          | 52 ++++++++++++++++++++++++++-
 2 files changed, 63 insertions(+), 12 deletions(-)

diff --git a/builder/digitalocean/step_snapshot.go b/builder/digitalocean/step_snapshot.go
index 8974ad23f..129a0d5b1 100644
--- a/builder/digitalocean/step_snapshot.go
+++ b/builder/digitalocean/step_snapshot.go
@@ -20,7 +20,7 @@ func (s *stepSnapshot) Run(state multistep.StateBag) multistep.StepAction {
 	dropletId := state.Get("droplet_id").(int)
 
 	ui.Say(fmt.Sprintf("Creating snapshot: %v", c.SnapshotName))
-	_, _, err := client.DropletActions.Snapshot(dropletId, c.SnapshotName)
+	action, _, err := client.DropletActions.Snapshot(dropletId, c.SnapshotName)
 	if err != nil {
 		err := fmt.Errorf("Error creating snapshot: %s", err)
 		state.Put("error", err)
@@ -28,22 +28,23 @@ func (s *stepSnapshot) Run(state multistep.StateBag) multistep.StepAction {
 		return multistep.ActionHalt
 	}
 
-	// Wait for the droplet to become unlocked first. For snapshots
-	// this can end up taking quite a long time, so we hardcode this to
-	// 20 minutes.
-	if err := waitForDropletUnlocked(client, dropletId, 20*time.Minute); err != nil {
+	// With the pending state over, verify that we're in the active state
+	ui.Say("Waiting for snapshot to complete...")
+	if err := waitForActionState(godo.ActionCompleted, dropletId, action.ID,
+		client, 20*time.Minute); err != nil {
 		// If we get an error the first time, actually report it
-		err := fmt.Errorf("Error shutting down droplet: %s", err)
+		err := fmt.Errorf("Error waiting for snapshot: %s", err)
 		state.Put("error", err)
 		ui.Error(err.Error())
 		return multistep.ActionHalt
 	}
 
-	// With the pending state over, verify that we're in the active state
-	ui.Say("Waiting for snapshot to complete...")
-	err = waitForDropletState("active", dropletId, client, c.StateTimeout)
-	if err != nil {
-		err := fmt.Errorf("Error waiting for snapshot to complete: %s", err)
+	// Wait for the droplet to become unlocked first. For snapshots
+	// this can end up taking quite a long time, so we hardcode this to
+	// 20 minutes.
+	if err := waitForDropletUnlocked(client, dropletId, 20*time.Minute); err != nil {
+		// If we get an error the first time, actually report it
+		err := fmt.Errorf("Error shutting down droplet: %s", err)
 		state.Put("error", err)
 		ui.Error(err.Error())
 		return multistep.ActionHalt
diff --git a/builder/digitalocean/wait.go b/builder/digitalocean/wait.go
index a41bbb3ed..35446674c 100644
--- a/builder/digitalocean/wait.go
+++ b/builder/digitalocean/wait.go
@@ -57,7 +57,7 @@ func waitForDropletUnlocked(
 	}
 }
 
-// waitForState simply blocks until the droplet is in
+// waitForDropletState simply blocks until the droplet is in
 // a state we expect, while eventually timing out.
 func waitForDropletState(
 	desiredState string, dropletId int,
@@ -106,3 +106,53 @@ func waitForDropletState(
 		return err
 	}
 }
+
+// waitForActionState simply blocks until the droplet action is in
+// a state we expect, while eventually timing out.
+func waitForActionState(
+	desiredState string, dropletId, actionId int,
+	client *godo.Client, timeout time.Duration) error {
+	done := make(chan struct{})
+	defer close(done)
+
+	result := make(chan error, 1)
+	go func() {
+		attempts := 0
+		for {
+			attempts += 1
+
+			log.Printf("Checking action status... (attempt: %d)", attempts)
+			action, _, err := client.DropletActions.Get(dropletId, actionId)
+			if err != nil {
+				result <- err
+				return
+			}
+
+			if action.Status == desiredState {
+				result <- nil
+				return
+			}
+
+			// Wait 3 seconds in between
+			time.Sleep(3 * time.Second)
+
+			// Verify we shouldn't exit
+			select {
+			case <-done:
+				// We finished, so just exit the goroutine
+				return
+			default:
+				// Keep going
+			}
+		}
+	}()
+
+	log.Printf("Waiting for up to %d seconds for action to become %s", timeout/time.Second, desiredState)
+	select {
+	case err := <-result:
+		return err
+	case <-time.After(timeout):
+		err := fmt.Errorf("Timeout while waiting to for action to become '%s'", desiredState)
+		return err
+	}
+}