diff --git a/state/state.go b/state/state.go index 39f81c9d1..246c9c44d 100644 --- a/state/state.go +++ b/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.