Merge pull request #861 from fraenkel/reclaim_lock

Reclaim locks

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)

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")
+	}
+}