diff --git a/server/mvcc/index.go b/server/mvcc/index.go index 57ba1bab4..0a5cb0051 100644 --- a/server/mvcc/index.go +++ b/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 }) diff --git a/server/mvcc/kv_test.go b/server/mvcc/kv_test.go index d7df00ebf..ad33b4041 100644 --- a/server/mvcc/kv_test.go +++ b/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) } } diff --git a/server/mvcc/kvstore_test.go b/server/mvcc/kvstore_test.go index 5de0a195c..1e56674d0 100644 --- a/server/mvcc/kvstore_test.go +++ b/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) } diff --git a/server/mvcc/kvstore_txn.go b/server/mvcc/kvstore_txn.go index 162228176..93d7db20e 100644 --- a/server/mvcc/kvstore_txn.go +++ b/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) { diff --git a/tests/integration/v3_grpc_test.go b/tests/integration/v3_grpc_test.go index 06d932197..920a759bd 100644 --- a/tests/integration/v3_grpc_test.go +++ b/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, 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)