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

Author: rjemanuele

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,30 @@ 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 path on timeout
+				delErr := l.c.Delete(path, -1)
+				if delErr != nil {
+					return delErr
+				}
+				return ErrTimeout
+			}
+		} else {
+			ev := <-ch
+			if ev.Err != nil {
+				return ev.Err
+			}
 		}
 	}
 

lock_test.go

@@ -1,6 +1,7 @@
 package zk
 
 import (
+	"errors"
 	"testing"
 	"time"
 )
@@ -61,6 +62,59 @@ func TestIntegration_Lock(t *testing.T) {
 	}
 }
 
+func TestIntegration_TryLock(t *testing.T) {
+	ts, err := StartTestCluster(t, 1, nil, logWriter{t: t, p: "[ZKERR] "})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ts.Stop()
+	zk, _, err := ts.ConnectAll()
+	if err != nil {
+		t.Fatalf("Connect returned error: %+v", err)
+	}
+	defer zk.Close()
+
+	acls := WorldACL(PermAll)
+
+	l := NewLock(zk, "/test", acls)
+	if err := l.TryLock(time.Second); err != nil {
+		t.Fatal(err)
+	}
+
+	l2 := NewLock(zk, "/test", acls)
+	if err := l2.TryLock(time.Millisecond * 100); !errors.Is(err, ErrTimeout) {
+		t.Fatalf("Expected ErrTimeout instead of %s", err)
+	}
+
+	val := make(chan int, 3)
+
+	go func() {
+		if err := l2.TryLock(time.Second); err != nil {
+			t.Fatal(err)
+		}
+		val <- 2
+		if err := l2.Unlock(); err != nil {
+			t.Fatal(err)
+		}
+		val <- 3
+	}()
+	time.Sleep(time.Millisecond * 100)
+
+	val <- 1
+	if err := l.Unlock(); err != nil {
+		t.Fatal(err)
+	}
+	if x := <-val; x != 1 {
+		t.Fatalf("Expected 1 instead of %d", x)
+	}
+	if x := <-val; x != 2 {
+		t.Fatalf("Expected 2 instead of %d", x)
+	}
+	if x := <-val; x != 3 {
+		t.Fatalf("Expected 3 instead of %d", x)
+	}
+}
+
 // This tests creating a lock with a path that's more than 1 node deep (e.g. "/test-multi-level/lock"),
 // when a part of that path already exists (i.e. "/test-multi-level" node already exists).
 func TestIntegration_MultiLevelLock(t *testing.T) {