lessor: fix le.itemMap leak after lease is revoked

Signed-off-by: qsyqian <[email protected]>
Co-authored-by: Thomas Jungblut <[email protected]>

server/lease/lessor.go

@@ -348,6 +348,9 @@ func (le *lessor) Revoke(id LeaseID) error {
 	sort.StringSlice(keys).Sort()
 	for _, key := range keys {
 		txn.DeleteRange([]byte(key), nil)
+		le.mu.Lock()
+		delete(le.itemMap, LeaseItem{Key: key})
+		le.mu.Unlock()
 	}
 
 	// lease deletion needs to be in the same backend transaction with the

server/lease/lessor_test.go

@@ -26,6 +26,7 @@ import (
 	"time"
 
 	"github.com/coreos/go-semver/semver"
+	"github.com/stretchr/testify/assert"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zaptest"
 
@@ -208,6 +209,39 @@ func TestLessorRevoke(t *testing.T) {
 	}
 }
 
+func TestLessorGrantAttachRevokeDetach(t *testing.T) {
+	lg := zap.NewNop()
+	_, be := NewTestBackend(t)
+	defer be.Close()
+
+	le := newLessor(lg, be, clusterLatest(), LessorConfig{MinLeaseTTL: minLeaseTTL})
+	defer le.Stop()
+	le.SetRangeDeleter(func() TxnDelete { return newFakeDeleter(be) })
+
+	// grant a lease with long term (100 seconds) to
+	// avoid early termination during the test.
+	l, err := le.Grant(1, 100)
+	if err != nil {
+		t.Fatalf("could not grant lease for 100s ttl (%v)", err)
+	}
+
+	items := []LeaseItem{
+		{"foo"},
+		{"bar"},
+	}
+
+	if err := le.Attach(l.ID, items); err != nil {
+		t.Fatalf("failed to attach items to the lease: %v", err)
+	}
+
+	if err := le.Revoke(l.ID); err != nil {
+		t.Fatalf("failed to revoke lease: %v", err)
+	}
+
+	assert.Equal(t, 0, len(le.leaseMap))
+	assert.Equal(t, 0, len(le.itemMap))
+}
+
 func renew(t *testing.T, le *lessor, id LeaseID) int64 {
 	ch := make(chan int64, 1)
 	errch := make(chan error, 1)

server/storage/mvcc/kvstore_test.go

@@ -29,6 +29,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"go.uber.org/zap/zaptest"
 
 	"go.etcd.io/etcd/api/v3/mvccpb"
@@ -178,6 +179,94 @@ func TestStorePut(t *testing.T) {
 	}
 }
 
+func TestStorePutDeleteWithLease(t *testing.T) {
+	lg := zaptest.NewLogger(t)
+	key := newTestKeyBytes(lg, revision{2, 0}, false)
+	kv := mvccpb.KeyValue{
+		Key:            []byte("foo"),
+		Value:          []byte("bar"),
+		CreateRevision: 2,
+		ModRevision:    2,
+		Version:        1,
+		Lease:          1,
+	}
+	kvb, err := kv.Marshal()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	s := newFakeStore(lg)
+	s.le = lease.NewLessor(lg, s.b, nil, lease.LessorConfig{MinLeaseTTL: 5})
+	b := s.b.(*fakeBackend)
+	fi := s.kvindex.(*fakeIndex)
+
+	l, err := s.le.Grant(1, 10)
+	if err != nil {
+		t.Errorf("failed to grant lease with ttl(10s), err: %v", err)
+	}
+
+	fi.indexGetRespc <- indexGetResp{revision{}, revision{}, 0, ErrRevisionNotFound}
+	s.Put([]byte("foo"), []byte("bar"), l.ID)
+
+	assert.Equal(t, []string{"foo"}, s.le.Lookup(l.ID).Keys())
+	assert.Equal(t, l.ID, s.le.GetLease(lease.LeaseItem{Key: "foo"}))
+
+	fi.indexRangeRespc <- indexRangeResp{[][]byte{[]byte("foo")}, []revision{{2, 0}}}
+	b.tx.rangeRespc <- rangeResp{[][]byte{key}, [][]byte{kvb}}
+	s.DeleteRange([]byte("foo"), nil)
+
+	assert.Equal(t, []string{}, s.le.Lookup(l.ID).Keys())
+	assert.Equal(t, lease.NoLease, s.le.GetLease(lease.LeaseItem{Key: "foo"}))
+
+	s.Close()
+}
+
+func TestStorePutWithLeaseRevoke(t *testing.T) {
+	lg := zaptest.NewLogger(t)
+	key := newTestKeyBytes(lg, revision{2, 0}, false)
+	kv := mvccpb.KeyValue{
+		Key:            []byte("foo"),
+		Value:          []byte("bar"),
+		CreateRevision: 2,
+		ModRevision:    2,
+		Version:        1,
+		Lease:          1,
+	}
+	kvb, err := kv.Marshal()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	s := newFakeStore(lg)
+	s.le = lease.NewLessor(lg, s.b, nil, lease.LessorConfig{MinLeaseTTL: 5})
+	s.le.SetRangeDeleter(func() lease.TxnDelete { return s.Write(traceutil.TODO()) })
+	b := s.b.(*fakeBackend)
+	fi := s.kvindex.(*fakeIndex)
+
+	l, err := s.le.Grant(1, 10)
+	if err != nil {
+		t.Errorf("failed to grant lease with ttl(10s), err: %v", err)
+	}
+
+	fi.indexGetRespc <- indexGetResp{revision{}, revision{}, 0, ErrRevisionNotFound}
+	s.Put([]byte("foo"), []byte("bar"), l.ID)
+
+	assert.Equal(t, []string{"foo"}, s.le.Lookup(l.ID).Keys())
+	assert.Equal(t, l.ID, s.le.GetLease(lease.LeaseItem{Key: "foo"}))
+
+	fi.indexRangeRespc <- indexRangeResp{[][]byte{[]byte("foo")}, []revision{{2, 0}}}
+	b.tx.rangeRespc <- rangeResp{[][]byte{key}, [][]byte{kvb}}
+	if err := s.le.Revoke(l.ID); err != nil {
+		t.Errorf("failed to revoke lease ID: %d, err: %v", l.ID, err)
+	}
+
+	assert.Nil(t, s.le.Lookup(l.ID))
+	assert.Equal(t, lease.NoLease, s.le.GetLease(lease.LeaseItem{Key: "foo"}))
+
+	s.Close()
+
+}
+
 func TestStoreRange(t *testing.T) {
 	lg := zaptest.NewLogger(t)
 	key := newTestKeyBytes(lg, revision{2, 0}, false)