From 9645107cf93108c512b189d0765bb9feb53b7221 Mon Sep 17 00:00:00 2001
From: John Engelman <john.r.engelman@gmail.com>
Date: Wed, 13 Jan 2016 16:37:57 -0600
Subject: [PATCH 1/2] Wait for upgrade to complete before confirming.

---
 main.go | 39 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 37 insertions(+), 2 deletions(-)

diff --git a/main.go b/main.go
index 8ce8875..645134b 100644
--- a/main.go
+++ b/main.go
@@ -2,8 +2,10 @@ package main
 
 import (
 	"fmt"
+	"log"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/drone/drone-go/plugin"
 	"github.com/rancher/go-rancher/client"
@@ -115,11 +117,44 @@ func main() {
 	fmt.Printf("Upgraded %s to %s\n", vargs.Service, vargs.Image)
 
 	if vargs.Confirm {
-		_, err = rancher.Service.ActionFinishupgrade(&service)
+		err = retryFunc(30*time.Second, func() error {
+			s, e := rancher.Service.ById(service.Id)
+			if e != nil {
+				return e
+			}
+			if s.State != "upgraded" {
+				return fmt.Errorf("Service not upgraded: %s", s.State)
+			}
+			return nil
+		})
 		if err != nil {
-			fmt.Printf("Unable to finish upgrade %s\n", vargs.Service)
+			fmt.Printf("Error waiting for service upgrade to complete: %s", err)
+			os.Exit(1)
+		}
+
+		s, err := rancher.Service.ById(service.Id)
+		_, err = rancher.Service.ActionFinishupgrade(s)
+		if err != nil {
+			fmt.Printf("Unable to finish upgrade %s: %s\n", vargs.Service, err)
 			os.Exit(1)
 		}
 		fmt.Printf("Finished upgrade %s\n", vargs.Service)
 	}
 }
+
+func retryFunc(timeout time.Duration, f func() error) error {
+	finish := time.After(timeout)
+	for {
+		err := f()
+		if err == nil {
+			return nil
+		}
+		log.Printf("Retryable error: %v", err)
+
+		select {
+		case <-finish:
+			return err
+		case <-time.After(3 * time.Second):
+		}
+	}
+}

From 1e22953e15b80c63968528215b91cf57e58e1ba9 Mon Sep 17 00:00:00 2001
From: John Engelman <john.r.engelman@gmail.com>
Date: Wed, 13 Jan 2016 17:59:32 -0600
Subject: [PATCH 2/2] Cleanup and retry more reusable.

---
 main.go | 29 ++++++++++++++---------------
 1 file changed, 14 insertions(+), 15 deletions(-)

diff --git a/main.go b/main.go
index 645134b..065aac0 100644
--- a/main.go
+++ b/main.go
@@ -2,7 +2,6 @@ package main
 
 import (
 	"fmt"
-	"log"
 	"os"
 	"strings"
 	"time"
@@ -117,23 +116,23 @@ func main() {
 	fmt.Printf("Upgraded %s to %s\n", vargs.Service, vargs.Image)
 
 	if vargs.Confirm {
-		err = retryFunc(30*time.Second, func() error {
+		srv, err := retry(func() (interface{}, error) {
 			s, e := rancher.Service.ById(service.Id)
 			if e != nil {
-				return e
+				return nil, e
 			}
 			if s.State != "upgraded" {
-				return fmt.Errorf("Service not upgraded: %s", s.State)
+				return nil, fmt.Errorf("Service not upgraded: %s", s.State)
 			}
-			return nil
-		})
+			return s, nil
+		}, 30*time.Second, 3*time.Second)
+
 		if err != nil {
 			fmt.Printf("Error waiting for service upgrade to complete: %s", err)
 			os.Exit(1)
 		}
 
-		s, err := rancher.Service.ById(service.Id)
-		_, err = rancher.Service.ActionFinishupgrade(s)
+		_, err = rancher.Service.ActionFinishupgrade(srv.(*client.Service))
 		if err != nil {
 			fmt.Printf("Unable to finish upgrade %s: %s\n", vargs.Service, err)
 			os.Exit(1)
@@ -142,19 +141,19 @@ func main() {
 	}
 }
 
-func retryFunc(timeout time.Duration, f func() error) error {
+type retryFunc func() (interface{}, error)
+
+func retry(f retryFunc, timeout time.Duration, interval time.Duration) (interface{}, error) {
 	finish := time.After(timeout)
 	for {
-		err := f()
+		result, err := f()
 		if err == nil {
-			return nil
+			return result, nil
 		}
-		log.Printf("Retryable error: %v", err)
-
 		select {
 		case <-finish:
-			return err
-		case <-time.After(3 * time.Second):
+			return nil, err
+		case <-time.After(interval):
 		}
 	}
 }