add state.LockWithContext

LockWithContext will retry a lock until the context expires or is
cancelled. This will let us implement a `-lock-timeout` flag, and make
use of existing contexts when applicable.

state/state.go

@@ -2,6 +2,7 @@ package state
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -73,6 +74,36 @@ type Locker interface {
 	Unlock(id string) error
 }
 
+// Lock the state, using the provided context for timeout and cancellation
+// TODO: this should probably backoff somewhat.
+func LockWithContext(s State, info *LockInfo, ctx context.Context) (string, error) {
+	for {
+		id, err := s.Lock(info)
+		if err == nil {
+			return id, nil
+		}
+
+		le, ok := err.(*LockError)
+		if !ok {
+			// not a lock error, so we can't retry
+			return "", err
+		}
+
+		if le.Info.ID == "" {
+			// the lock has no ID, something is wrong so don't keep trying
+			return "", fmt.Errorf("lock error missing ID: %s", err)
+		}
+
+		// there's an existing lock, wait and try again
+		select {
+		case <-ctx.Done():
+			// return the last lock error with the info
+			return "", err
+		case <-time.After(time.Second):
+		}
+	}
+}
+
 // Generate a LockInfo structure, populating the required fields.
 func NewLockInfo() *LockInfo {
 	// this doesn't need to be cryptographically secure, just unique.