From 717746ce4ff3b1e2b5d282c028094b7ff45fa59c Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Mon, 19 Aug 2013 23:02:06 -0700
Subject: [PATCH] provisioner/shell: retry uploads if reboot [GH-282]

---
 CHANGELOG.md                     |  5 +++
 provisioner/shell/provisioner.go | 53 ++++++++++++++++++++------------
 2 files changed, 39 insertions(+), 19 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 163408e4e..c505f8763 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,11 @@ IMPROVEMENTS:
 * post-processor/vagrant: the file being compressed will be shown
   in the UI [GH-314]
 
+BUG FIXES:
+
+* provisioner/shell: Retry shell script uploads, making reboots more
+  robust if they happen to fail in this stage. [GH-282]
+
 ## 0.3.3 (August 19, 2013)
 
 FEATURES:
diff --git a/provisioner/shell/provisioner.go b/provisioner/shell/provisioner.go
index b84461dfc..05580aca6 100644
--- a/provisioner/shell/provisioner.go
+++ b/provisioner/shell/provisioner.go
@@ -237,7 +237,9 @@ func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
 		defer f.Close()
 
 		log.Printf("Uploading %s => %s", path, p.config.RemotePath)
-		err = comm.Upload(p.config.RemotePath, f)
+		err = p.retryable(func() error {
+			return comm.Upload(p.config.RemotePath, f)
+		})
 		if err != nil {
 			return fmt.Errorf("Error uploading shell script: %s", err)
 		}
@@ -258,26 +260,13 @@ func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
 		}
 
 		cmd := &packer.RemoteCmd{Command: command}
-		startTimeout := time.After(p.config.startRetryTimeout)
 		log.Printf("Executing command: %s", cmd.Command)
-		for {
-			if err := cmd.StartWithUi(comm, ui); err == nil {
-				break
-			}
+		err = p.retryable(func() error {
+			return cmd.StartWithUi(comm, ui)
+		})
 
-			// Create an error and log it
-			err = fmt.Errorf("Error executing command: %s", err)
-			log.Printf(err.Error())
-
-			// Check if we timed out, otherwise we retry. It is safe to
-			// retry since the only error case above is if the command
-			// failed to START.
-			select {
-			case <-startTimeout:
-				return err
-			default:
-				time.Sleep(2 * time.Second)
-			}
+		if err != nil {
+			return err
 		}
 
 		if cmd.ExitStatus != 0 {
@@ -287,3 +276,29 @@ func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
 
 	return nil
 }
+
+// retryable will retry the given function over and over until a
+// non-error is returned.
+func (p *Provisioner) retryable(f func() error) error {
+	startTimeout := time.After(p.config.startRetryTimeout)
+	for {
+		var err error
+		if err = f(); err == nil {
+			return nil
+		}
+
+		// Create an error and log it
+		err = fmt.Errorf("Retryable error: %s", err)
+		log.Printf(err.Error())
+
+		// Check if we timed out, otherwise we retry. It is safe to
+		// retry since the only error case above is if the command
+		// failed to START.
+		select {
+		case <-startTimeout:
+			return err
+		default:
+			time.Sleep(2 * time.Second)
+		}
+	}
+}