From 093b23d0c320b1c2d9cf3a9dfbfcabf97f2d050e Mon Sep 17 00:00:00 2001
From: Michael Fraenkel <fraenkel@us.ibm.com>
Date: Sun, 12 Apr 2015 16:48:11 -0500
Subject: [PATCH] Reclaim locks

- When attempting to lock with the same session that already owns the lock,
  Lock() will re-acquire the lock.
---
 api/lock.go      |  9 +++++--
 api/lock_test.go | 67 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 2 deletions(-)

diff --git a/api/lock.go b/api/lock.go
index f6fdbbb16..4b694789c 100644
--- a/api/lock.go
+++ b/api/lock.go
@@ -165,14 +165,18 @@ WAIT:
 	if pair != nil && pair.Flags != LockFlagValue {
 		return nil, ErrLockConflict
 	}
+	locked := false
+	if pair != nil && pair.Session == l.lockSession {
+		goto HELD
+	}
 	if pair != nil && pair.Session != "" {
 		qOpts.WaitIndex = meta.LastIndex
 		goto WAIT
 	}
 
 	// Try to acquire the lock
-	lockEnt := l.lockEntry(l.lockSession)
-	locked, _, err := kv.Acquire(lockEnt, nil)
+	pair = l.lockEntry(l.lockSession)
+	locked, _, err = kv.Acquire(pair, nil)
 	if err != nil {
 		return nil, fmt.Errorf("failed to acquire lock: %v", err)
 	}
@@ -187,6 +191,7 @@ WAIT:
 		}
 	}
 
+HELD:
 	// Watch to ensure we maintain leadership
 	leaderCh := make(chan struct{})
 	go l.monitorLock(l.lockSession, leaderCh)
diff --git a/api/lock_test.go b/api/lock_test.go
index a4aea7349..de382f4b9 100644
--- a/api/lock_test.go
+++ b/api/lock_test.go
@@ -287,3 +287,70 @@ func TestLock_Conflict(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	}
 }
+
+func TestLock_ReclaimLock(t *testing.T) {
+	c, s := makeClient(t)
+	defer s.stop()
+
+	session, _, err := c.Session().Create(&SessionEntry{}, nil)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	lock, err := c.LockOpts(&LockOptions{Key: "test/lock", Session: session})
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Should work
+	leaderCh, err := lock.Lock(nil)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if leaderCh == nil {
+		t.Fatalf("not leader")
+	}
+	defer lock.Unlock()
+
+	l2, err := c.LockOpts(&LockOptions{Key: "test/lock", Session: session})
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	reclaimed := make(chan (<-chan struct{}), 1)
+	go func() {
+		l2Ch, err := l2.Lock(nil)
+		if err != nil {
+			t.Fatalf("not locked: %v", err)
+		}
+		reclaimed <- l2Ch
+	}()
+
+	// Should reclaim the lock
+	var leader2Ch <-chan struct{}
+
+	select {
+	case leader2Ch = <-reclaimed:
+	case <-time.After(time.Second):
+		t.Fatalf("should have locked")
+	}
+
+	// unlock should work
+	err = l2.Unlock()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	//Both locks should see the unlock
+	select {
+	case <-leader2Ch:
+	case <-time.After(time.Second):
+		t.Fatalf("should not be leader")
+	}
+
+	select {
+	case <-leaderCh:
+	case <-time.After(time.Second):
+		t.Fatalf("should not be leader")
+	}
+}