From f0c184b3a23af4b3cb40522e42fdfeb978dfb991 Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Wed, 4 Jan 2017 22:30:03 -0800
Subject: [PATCH] lease: support mvcc txn

---
 lease/lessor.go      | 45 +++++++++++++++++---------------------------
 lease/lessor_test.go | 22 ++++++----------------
 2 files changed, 23 insertions(+), 44 deletions(-)

diff --git a/lease/lessor.go b/lease/lessor.go
index 385bd76d7..5120d1cfc 100644
--- a/lease/lessor.go
+++ b/lease/lessor.go
@@ -43,28 +43,24 @@ var (
 	ErrLeaseExists   = errors.New("lease already exists")
 )
 
-type LeaseID int64
-
-// RangeDeleter defines an interface with Txn and DeleteRange method.
-// We define this interface only for lessor to limit the number
-// of methods of mvcc.KV to what lessor actually needs.
-//
-// Having a minimum interface makes testing easy.
-type RangeDeleter interface {
-	// TxnBegin see comments on mvcc.KV
-	TxnBegin() int64
-	// TxnEnd see comments on mvcc.KV
-	TxnEnd(txnID int64) error
-	// TxnDeleteRange see comments on mvcc.KV
-	TxnDeleteRange(txnID int64, key, end []byte) (n, rev int64, err error)
+// TxnDelete is a TxnWrite that only permits deletes. Defined here
+// to avoid circular dependency with mvcc.
+type TxnDelete interface {
+	DeleteRange(key, end []byte) (n, rev int64)
+	End()
 }
 
+// RangeDeleter is a TxnDelete constructor.
+type RangeDeleter func() TxnDelete
+
+type LeaseID int64
+
 // Lessor owns leases. It can grant, revoke, renew and modify leases for lessee.
 type Lessor interface {
-	// SetRangeDeleter sets the RangeDeleter to the Lessor.
-	// Lessor deletes the items in the revoked or expired lease from the
-	// the set RangeDeleter.
-	SetRangeDeleter(dr RangeDeleter)
+	// SetRangeDeleter lets the lessor create TxnDeletes to the store.
+	// Lessor deletes the items in the revoked or expired lease by creating
+	// new TxnDeletes.
+	SetRangeDeleter(rd RangeDeleter)
 
 	// Grant grants a lease that expires at least after TTL seconds.
 	Grant(id LeaseID, ttl int64) (*Lease, error)
@@ -248,17 +244,14 @@ func (le *lessor) Revoke(id LeaseID) error {
 		return nil
 	}
 
-	tid := le.rd.TxnBegin()
+	txn := le.rd()
 
 	// sort keys so deletes are in same order among all members,
 	// otherwise the backened hashes will be different
 	keys := l.Keys()
 	sort.StringSlice(keys).Sort()
 	for _, key := range keys {
-		_, _, err := le.rd.TxnDeleteRange(tid, []byte(key), nil)
-		if err != nil {
-			panic(err)
-		}
+		txn.DeleteRange([]byte(key), nil)
 	}
 
 	le.mu.Lock()
@@ -269,11 +262,7 @@ func (le *lessor) Revoke(id LeaseID) error {
 	// deleting the keys if etcdserver fails in between.
 	le.b.BatchTx().UnsafeDelete(leaseBucketName, int64ToBytes(int64(l.ID)))
 
-	err := le.rd.TxnEnd(tid)
-	if err != nil {
-		panic(err)
-	}
-
+	txn.End()
 	return nil
 }
 
diff --git a/lease/lessor_test.go b/lease/lessor_test.go
index 97793f694..c90b64aba 100644
--- a/lease/lessor_test.go
+++ b/lease/lessor_test.go
@@ -86,10 +86,8 @@ func TestLeaseConcurrentKeys(t *testing.T) {
 	defer os.RemoveAll(dir)
 	defer be.Close()
 
-	fd := &fakeDeleter{}
-
 	le := newLessor(be, minLeaseTTL)
-	le.SetRangeDeleter(fd)
+	le.SetRangeDeleter(func() TxnDelete { return &fakeDeleter{} })
 
 	// grant a lease with long term (100 seconds) to
 	// avoid early termination during the test.
@@ -138,7 +136,7 @@ func TestLessorRevoke(t *testing.T) {
 	fd := &fakeDeleter{}
 
 	le := newLessor(be, minLeaseTTL)
-	le.SetRangeDeleter(fd)
+	le.SetRangeDeleter(func() TxnDelete { return fd })
 
 	// grant a lease with long term (100 seconds) to
 	// avoid early termination during the test.
@@ -215,10 +213,8 @@ func TestLessorDetach(t *testing.T) {
 	defer os.RemoveAll(dir)
 	defer be.Close()
 
-	fd := &fakeDeleter{}
-
 	le := newLessor(be, minLeaseTTL)
-	le.SetRangeDeleter(fd)
+	le.SetRangeDeleter(func() TxnDelete { return &fakeDeleter{} })
 
 	// grant a lease with long term (100 seconds) to
 	// avoid early termination during the test.
@@ -382,17 +378,11 @@ type fakeDeleter struct {
 	deleted []string
 }
 
-func (fd *fakeDeleter) TxnBegin() int64 {
-	return 0
-}
+func (fd *fakeDeleter) End() {}
 
-func (fd *fakeDeleter) TxnEnd(txnID int64) error {
-	return nil
-}
-
-func (fd *fakeDeleter) TxnDeleteRange(tid int64, key, end []byte) (int64, int64, error) {
+func (fd *fakeDeleter) DeleteRange(key, end []byte) (int64, int64) {
 	fd.deleted = append(fd.deleted, string(key)+"_"+string(end))
-	return 0, 0, nil
+	return 0, 0
 }
 
 func NewTestBackend(t *testing.T) (string, backend.Backend) {