Add in support for a TryLock to attempt to get a lock within a timeout

constants.go

@@ -124,6 +124,7 @@ var (
 	ErrSessionMoved            = errors.New("zk: session moved to another server, so operation is ignored")
 	ErrReconfigDisabled        = errors.New("attempts to perform a reconfiguration operation when reconfiguration feature is disabled")
 	ErrBadArguments            = errors.New("invalid arguments")
+	ErrTimeout                 = errors.New("timeout exceeded")
 	// ErrInvalidCallback         = errors.New("zk: invalid callback specified")
 
 	errCodeToError = map[ErrCode]error{
@@ -144,6 +145,7 @@ var (
 		errSessionMoved:      ErrSessionMoved,
 		errZReconfigDisabled: ErrReconfigDisabled,
 		errBadArguments:      ErrBadArguments,
+		errTimeout:           ErrTimeout,
 	}
 )
 
@@ -166,6 +168,7 @@ const (
 	errOperationTimeout     = -7
 	errBadArguments         = -8
 	errInvalidState         = -9
+	errTimeout              = -10
 	// API errors
 	errAPIError                ErrCode = -100
 	errNoNode                  ErrCode = -101 // *

lock.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"time"
 )
 
 var (
@@ -49,10 +50,24 @@ func (l *Lock) Lock() error {
 	return l.LockWithData([]byte{})
 }
 
+// TryLock attempts to acquire the lock before a timeout. It works like LockWithData, but it doesn't
+// write any data to the lock node.
+func (l *Lock) TryLock(timeout time.Duration) error {
+	return l.TryLockWithData([]byte{}, timeout)
+}
+
 // LockWithData attempts to acquire the lock, writing data into the lock node.
 // It will wait to return until the lock is acquired or an error occurs. If
 // this instance already has the lock then ErrDeadlock is returned.
 func (l *Lock) LockWithData(data []byte) error {
+	return l.TryLockWithData(data, time.Duration(-1))
+}
+
+// TryLockWithData attempts to acquire the lock, writing data into the lock node.
+// It will wait to return until the lock is acquired, an error occurs, or the timeout. If
+// this instance already has the lock then ErrDeadlock is returned.  A negative
+// timeout is waiting forever
+func (l *Lock) TryLockWithData(data []byte, timeout time.Duration) error {
 	if l.lockPath != "" {
 		return ErrDeadlock
 	}
@@ -97,6 +112,8 @@ func (l *Lock) LockWithData(data []byte) error {
 		return err
 	}
 
+	start := time.Now()
+
 	for {
 		children, _, err := l.c.Children(l.path)
 		if err != nil {
@@ -134,9 +151,27 @@ func (l *Lock) LockWithData(data []byte) error {
 			continue
 		}
 
-		ev := <-ch
-		if ev.Err != nil {
-			return ev.Err
+		if timeout >= 0 {
+			wait := start.Add(timeout).Sub(time.Now())
+			if wait < 0 {
+				wait = time.Duration(1)
+			}
+			delay := time.NewTimer(wait)
+			select {
+			case ev := <-ch:
+				if ev.Err != nil {
+					return ev.Err
+				}
+			case <-delay.C:
+				// remove the pending lock on timeout
+				_ = l.Unlock()
+				return ErrTimeout
+			}
+		} else {
+			ev := <-ch
+			if ev.Err != nil {
+				return ev.Err
+			}
 		}
 	}