From abf2e618b5198bdf23f1cd41d35da0b5a5d203ab Mon Sep 17 00:00:00 2001
From: Luke Farnell <lfarnell@outlook.com>
Date: Tue, 16 May 2017 10:38:53 -0400
Subject: [PATCH] Wait for snapshot transfer and change artifact output

---
 builder/digitalocean/artifact.go      |  7 ++--
 builder/digitalocean/builder.go       |  2 +-
 builder/digitalocean/step_snapshot.go | 10 +++++-
 builder/digitalocean/wait.go          | 50 +++++++++++++++++++++++++++
 4 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/builder/digitalocean/artifact.go b/builder/digitalocean/artifact.go
index 9630686e3..54aed4786 100644
--- a/builder/digitalocean/artifact.go
+++ b/builder/digitalocean/artifact.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"strconv"
+	"strings"
 
 	"github.com/digitalocean/godo"
 )
@@ -17,7 +18,7 @@ type Artifact struct {
 	snapshotId int
 
 	// The name of the region
-	regionName string
+	regionNames []string
 
 	// The client for making API calls
 	client *godo.Client
@@ -33,11 +34,11 @@ func (*Artifact) Files() []string {
 }
 
 func (a *Artifact) Id() string {
-	return fmt.Sprintf("%s:%s", a.regionName, strconv.FormatUint(uint64(a.snapshotId), 10))
+	return fmt.Sprintf("%s:%s", strings.Join(a.regionNames[:], ","), strconv.FormatUint(uint64(a.snapshotId), 10))
 }
 
 func (a *Artifact) String() string {
-	return fmt.Sprintf("A snapshot was created: '%v' (ID: %v) in region '%v'", a.snapshotName, a.snapshotId, a.regionName)
+	return fmt.Sprintf("A snapshot was created: '%v' (ID: %v) in regions '%v'", a.snapshotName, a.snapshotId, strings.Join(a.regionNames[:], ","))
 }
 
 func (a *Artifact) State(name string) interface{} {
diff --git a/builder/digitalocean/builder.go b/builder/digitalocean/builder.go
index 1dc93b793..1c4572058 100644
--- a/builder/digitalocean/builder.go
+++ b/builder/digitalocean/builder.go
@@ -89,7 +89,7 @@ func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packe
 	artifact := &Artifact{
 		snapshotName: state.Get("snapshot_name").(string),
 		snapshotId:   state.Get("snapshot_image_id").(int),
-		regionName:   state.Get("region").(string),
+		regionNames:  state.Get("regions").([]string),
 		client:       client,
 	}
 
diff --git a/builder/digitalocean/step_snapshot.go b/builder/digitalocean/step_snapshot.go
index 6ec3e4245..180cc8fcb 100644
--- a/builder/digitalocean/step_snapshot.go
+++ b/builder/digitalocean/step_snapshot.go
@@ -90,6 +90,14 @@ func (s *stepSnapshot) Run(state multistep.StateBag) multistep.StepAction {
 				return multistep.ActionHalt
 			}
 			ui.Say(fmt.Sprintf("Transfering Snapshot ID: %d", imageTransfer.ID))
+			if err := waitForImageState(godo.ActionCompleted, imageTransfer.ID, action.ID,
+				client, 20*time.Minute); err != nil {
+				// If we get an error the first time, actually report it
+				err := fmt.Errorf("Error waiting for snapshot transfer: %s", err)
+				state.Put("error", err)
+				ui.Error(err.Error())
+				return multistep.ActionHalt
+			}
 		}
 	}
 
@@ -106,7 +114,7 @@ func (s *stepSnapshot) Run(state multistep.StateBag) multistep.StepAction {
 	log.Printf("Snapshot image ID: %d", imageId)
 	state.Put("snapshot_image_id", imageId)
 	state.Put("snapshot_name", c.SnapshotName)
-	state.Put("region", c.Region)
+	state.Put("regions", c.SnapshotRegions)
 
 	return multistep.ActionContinue
 }
diff --git a/builder/digitalocean/wait.go b/builder/digitalocean/wait.go
index e870096cd..2e4bd280c 100644
--- a/builder/digitalocean/wait.go
+++ b/builder/digitalocean/wait.go
@@ -157,3 +157,53 @@ func waitForActionState(
 		return err
 	}
 }
+
+// waitForImageState simply blocks until the image action is in
+// a state we expect, while eventually timing out.
+func waitForImageState(
+	desiredState string, imageId, 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.ImageActions.Get(context.TODO(), imageId, 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 image transter 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 image transter to become '%s'", desiredState)
+		return err
+	}
+}