Reworked WaitForHeight to avoid race conditions

rpc/client/helpers.go

@@ -6,26 +6,44 @@ import (
 	"github.com/pkg/errors"
 )
 
+// Waiter is informed of current height, decided whether to quit early
+type Waiter func(delta int) (abort error)
+
+// DefaultWaitStrategy is the standard backoff algorithm,
+// but you can plug in another one
+func DefaultWaitStrategy(delta int) (abort error) {
+	if delta > 10 {
+		return errors.Errorf("Waiting for %d blocks... aborting", delta)
+	} else if delta > 0 {
+		// estimate of wait time....
+		// wait half a second for the next block (in progress)
+		// plus one second for every full block
+		delay := time.Duration(delta-1)*time.Second + 500*time.Millisecond
+		time.Sleep(delay)
+	}
+	return nil
+}
+
 // Wait for height will poll status at reasonable intervals until
 // the block at the given height is available.
-func WaitForHeight(c StatusClient, h int) error {
-	wait := 1
-	for wait > 0 {
+//
+// If waiter is nil, we use DefaultWaitStrategy, but you can also
+// provide your own implementation
+func WaitForHeight(c StatusClient, h int, waiter Waiter) error {
+	if waiter == nil {
+		waiter = DefaultWaitStrategy
+	}
+	delta := 1
+	for delta > 0 {
 		s, err := c.Status()
 		if err != nil {
 			return err
 		}
-		wait = h - s.LatestBlockHeight
-		if wait > 10 {
-			return errors.Errorf("Waiting for %d block... aborting", wait)
-		} else if wait > 0 {
-			// estimate of wait time....
-			// wait half a second for the next block (in progress)
-			// plus one second for every full block
-			delay := time.Duration(wait-1)*time.Second + 500*time.Millisecond
-			time.Sleep(delay)
+		delta = h - s.LatestBlockHeight
+		// wait for the time, or abort early
+		if err := waiter(delta); err != nil {
+			return err
 		}
 	}
-	// guess we waited long enough
 	return nil
 }