etcdserver: Fix invalid count returned on Range with Limit

server/mvcc/index.go

@@ -25,8 +25,8 @@ import (
 type index interface {
 	Get(key []byte, atRev int64) (rev, created revision, ver int64, err error)
 	Range(key, end []byte, atRev int64) ([][]byte, []revision)
-	Revisions(key, end []byte, atRev int64, limit int) []revision
-	CountRevisions(key, end []byte, atRev int64, limit int) int
+	Revisions(key, end []byte, atRev int64, limit int) ([]revision, int)
+	CountRevisions(key, end []byte, atRev int64) int
 	Put(key []byte, rev revision)
 	Tombstone(key []byte, rev revision) error
 	RangeSince(key, end []byte, rev int64) []revision
@@ -106,27 +106,27 @@ func (ti *treeIndex) visit(key, end []byte, f func(ki *keyIndex) bool) {
 	})
 }
 
-func (ti *treeIndex) Revisions(key, end []byte, atRev int64, limit int) (revs []revision) {
+func (ti *treeIndex) Revisions(key, end []byte, atRev int64, limit int) (revs []revision, total int) {
 	if end == nil {
 		rev, _, _, err := ti.Get(key, atRev)
 		if err != nil {
-			return nil
+			return nil, 0
 		}
-		return []revision{rev}
+		return []revision{rev}, 1
 	}
 	ti.visit(key, end, func(ki *keyIndex) bool {
 		if rev, _, _, err := ki.get(ti.lg, atRev); err == nil {
-			revs = append(revs, rev)
-			if len(revs) == limit {
-				return false
+			if limit <= 0 || len(revs) < limit {
+				revs = append(revs, rev)
 			}
+			total++
 		}
 		return true
 	})
-	return revs
+	return revs, total
 }
 
-func (ti *treeIndex) CountRevisions(key, end []byte, atRev int64, limit int) int {
+func (ti *treeIndex) CountRevisions(key, end []byte, atRev int64) int {
 	if end == nil {
 		_, _, _, err := ti.Get(key, atRev)
 		if err != nil {
@@ -138,9 +138,6 @@ func (ti *treeIndex) CountRevisions(key, end []byte, atRev int64, limit int) int
 	ti.visit(key, end, func(ki *keyIndex) bool {
 		if _, _, _, err := ki.get(ti.lg, atRev); err == nil {
 			total++
-			if total == limit {
-				return false
-			}
 		}
 		return true
 	})

server/mvcc/kv_test.go

@@ -221,17 +221,18 @@ func testKVRangeLimit(t *testing.T, f rangeFunc) {
 
 	wrev := int64(4)
 	tests := []struct {
-		limit int64
-		wkvs  []mvccpb.KeyValue
+		limit   int64
+		wcounts int64
+		wkvs    []mvccpb.KeyValue
 	}{
 		// no limit
-		{-1, kvs},
+		{-1, 3, kvs},
 		// no limit
-		{0, kvs},
-		{1, kvs[:1]},
-		{2, kvs[:2]},
-		{3, kvs},
-		{100, kvs},
+		{0, 3, kvs},
+		{1, 3, kvs[:1]},
+		{2, 3, kvs[:2]},
+		{3, 3, kvs},
+		{100, 3, kvs},
 	}
 	for i, tt := range tests {
 		r, err := f(s, []byte("foo"), []byte("foo3"), RangeOptions{Limit: tt.limit})
@@ -248,7 +249,7 @@ func testKVRangeLimit(t *testing.T, f rangeFunc) {
 			if r.Count != len(kvs) {
 				t.Errorf("#%d: count = %d, want %d", i, r.Count, len(kvs))
 			}
-		} else if r.Count != int(tt.limit) {
+		} else if r.Count != int(tt.wcounts) {
 			t.Errorf("#%d: count = %d, want %d", i, r.Count, tt.limit)
 		}
 	}

server/mvcc/kvstore_test.go

@@ -937,12 +937,15 @@ type fakeIndex struct {
 	indexCompactRespc     chan map[revision]struct{}
 }
 
-func (i *fakeIndex) Revisions(key, end []byte, atRev int64, limit int) []revision {
+func (i *fakeIndex) Revisions(key, end []byte, atRev int64, limit int) ([]revision, int) {
 	_, rev := i.Range(key, end, atRev)
-	return rev
+	if len(rev) >= limit {
+		rev = rev[:limit]
+	}
+	return rev, len(rev)
 }
 
-func (i *fakeIndex) CountRevisions(key, end []byte, atRev int64, limit int) int {
+func (i *fakeIndex) CountRevisions(key, end []byte, atRev int64) int {
 	_, rev := i.Range(key, end, atRev)
 	return len(rev)
 }

server/mvcc/kvstore_txn.go

@@ -136,14 +136,14 @@ func (tr *storeTxnRead) rangeKeys(ctx context.Context, key, end []byte, curRev i
 		return &RangeResult{KVs: nil, Count: -1, Rev: 0}, ErrCompacted
 	}
 	if ro.Count {
-		total := tr.s.kvindex.CountRevisions(key, end, rev, int(ro.Limit))
+		total := tr.s.kvindex.CountRevisions(key, end, rev)
 		tr.trace.Step("count revisions from in-memory index tree")
 		return &RangeResult{KVs: nil, Count: total, Rev: curRev}, nil
 	}
-	revpairs := tr.s.kvindex.Revisions(key, end, rev, int(ro.Limit))
+	revpairs, total := tr.s.kvindex.Revisions(key, end, rev, int(ro.Limit))
 	tr.trace.Step("range keys from in-memory index tree")
 	if len(revpairs) == 0 {
-		return &RangeResult{KVs: nil, Count: 0, Rev: curRev}, nil
+		return &RangeResult{KVs: nil, Count: total, Rev: curRev}, nil
 	}
 
 	limit := int(ro.Limit)
@@ -176,7 +176,7 @@ func (tr *storeTxnRead) rangeKeys(ctx context.Context, key, end []byte, curRev i
 		}
 	}
 	tr.trace.Step("range keys from bolt db")
-	return &RangeResult{KVs: kvs, Count: len(revpairs), Rev: curRev}, nil
+	return &RangeResult{KVs: kvs, Count: total, Rev: curRev}, nil
 }
 
 func (tw *storeTxnWrite) put(key, value []byte, leaseID lease.LeaseID) {

tests/integration/v3_grpc_test.go

@@ -1289,8 +1289,9 @@ func TestV3RangeRequest(t *testing.T) {
 		putKeys []string
 		reqs    []pb.RangeRequest
 
-		wresps [][]string
-		wmores []bool
+		wresps  [][]string
+		wmores  []bool
+		wcounts []int64
 	}{
 		// single key
 		{
@@ -1307,6 +1308,7 @@ func TestV3RangeRequest(t *testing.T) {
 				{},
 			},
 			[]bool{false, false},
+			[]int64{1, 0},
 		},
 		// multi-key
 		{
@@ -1335,6 +1337,7 @@ func TestV3RangeRequest(t *testing.T) {
 				{"a", "b", "c", "d", "e"},
 			},
 			[]bool{false, false, false, false, false, false},
+			[]int64{5, 2, 0, 0, 0, 5},
 		},
 		// revision
 		{
@@ -1353,22 +1356,30 @@ func TestV3RangeRequest(t *testing.T) {
 				{"a", "b"},
 			},
 			[]bool{false, false, false, false},
+			[]int64{5, 0, 1, 2},
 		},
 		// limit
 		{
-			[]string{"foo", "bar"},
+			[]string{"a", "b", "c"},
 			[]pb.RangeRequest{
 				// more
 				{Key: []byte("a"), RangeEnd: []byte("z"), Limit: 1},
-				// no more
+				// half
 				{Key: []byte("a"), RangeEnd: []byte("z"), Limit: 2},
+				// no more
+				{Key: []byte("a"), RangeEnd: []byte("z"), Limit: 3},
+				// limit over
+				{Key: []byte("a"), RangeEnd: []byte("z"), Limit: 4},
 			},
 
 			[][]string{
-				{"bar"},
-				{"bar", "foo"},
+				{"a"},
+				{"a", "b"},
+				{"a", "b", "c"},
+				{"a", "b", "c"},
 			},
-			[]bool{true, false},
+			[]bool{true, true, false},
+			[]int64{3, 3, 3, 3},
 		},
 		// sort
 		{
@@ -1421,6 +1432,7 @@ func TestV3RangeRequest(t *testing.T) {
 				{"b", "a", "c", "d"},
 			},
 			[]bool{true, true, true, true, false, false},
+			[]int64{4, 4, 4, 4, 0, 4},
 		},
 		// min/max mod rev
 		{
@@ -1452,6 +1464,7 @@ func TestV3RangeRequest(t *testing.T) {
 				{"rev2", "rev3", "rev4", "rev5", "rev6"},
 			},
 			[]bool{false, false, false, false},
+			[]int64{5, 5, 5, 5},
 		},
 		// min/max create rev
 		{
@@ -1483,6 +1496,7 @@ func TestV3RangeRequest(t *testing.T) {
 				{"rev2", "rev3", "rev6"},
 			},
 			[]bool{false, false, false, false},
+			[]int64{3, 3, 3, 3},
 		},
 	}
 
@@ -1516,6 +1530,9 @@ func TestV3RangeRequest(t *testing.T) {
 			if resp.More != tt.wmores[j] {
 				t.Errorf("#%d.%d: bad more. got = %v, want = %v, ", i, j, resp.More, tt.wmores[j])
 			}
+			if resp.GetCount() != tt.wcounts[j] {
+				t.Errorf("#%d.%d: bad count. got = %v, want = %v, ", i, j, resp.GetCount(), tt.wcounts[j])
+			}
 			wrev := int64(len(tt.putKeys) + 1)
 			if resp.Header.Revision != wrev {
 				t.Errorf("#%d.%d: bad header revision. got = %d. want = %d", i, j, resp.Header.Revision, wrev)