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
 }

rpc/client/helpers_test.go

@@ -4,7 +4,6 @@ import (
 	"errors"
 	"strings"
 	"testing"
-	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -25,7 +24,7 @@ func TestWaitForHeight(t *testing.T) {
 	r := mock.NewStatusRecorder(m)
 
 	// connection failure always leads to error
-	err := client.WaitForHeight(r, 8)
+	err := client.WaitForHeight(r, 8, nil)
 	require.NotNil(err)
 	require.Equal("bye", err.Error())
 	// we called status once to check
@@ -37,26 +36,28 @@ func TestWaitForHeight(t *testing.T) {
 	}
 
 	// we will not wait for more than 10 blocks
-	err = client.WaitForHeight(r, 40)
+	err = client.WaitForHeight(r, 40, nil)
 	require.NotNil(err)
 	require.True(strings.Contains(err.Error(), "aborting"))
 	// we called status once more to check
 	require.Equal(2, len(r.Calls))
 
 	// waiting for the past returns immediately
-	err = client.WaitForHeight(r, 5)
+	err = client.WaitForHeight(r, 5, nil)
 	require.Nil(err)
 	// we called status once more to check
 	require.Equal(3, len(r.Calls))
 
-	// next tick in background while we wait
-	go func() {
-		time.Sleep(500 * time.Millisecond)
+	// since we can't update in a background goroutine (test --race)
+	// we use the callback to update the status height
+	myWaiter := func(delta int) error {
+		// update the height for the next call
 		m.Call.Response = &ctypes.ResultStatus{LatestBlockHeight: 15}
-	}()
+		return client.DefaultWaitStrategy(delta)
+	}
 
 	// we wait for a few blocks
-	err = client.WaitForHeight(r, 12)
+	err = client.WaitForHeight(r, 12, myWaiter)
 	require.Nil(err)
 	// we called status once to check
 	require.Equal(5, len(r.Calls))