diff --git a/integration/v3_grpc_test.go b/integration/v3_grpc_test.go index 331f2bc60..ce5dc379c 100644 --- a/integration/v3_grpc_test.go +++ b/integration/v3_grpc_test.go @@ -1285,8 +1285,9 @@ func TestV3RangeRequest(t *testing.T) { putKeys []string reqs []pb.RangeRequest - wresps [][]string - wmores []bool + wresps [][]string + wmores []bool + wcounts []int64 }{ // single key { @@ -1303,6 +1304,7 @@ func TestV3RangeRequest(t *testing.T) { {}, }, []bool{false, false}, + []int64{1, 0}, }, // multi-key { @@ -1331,6 +1333,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 { @@ -1349,22 +1352,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 { @@ -1417,6 +1428,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 { @@ -1448,6 +1460,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 { @@ -1479,6 +1492,7 @@ func TestV3RangeRequest(t *testing.T) { {"rev2", "rev3", "rev6"}, }, []bool{false, false, false, false}, + []int64{3, 3, 3, 3}, }, } @@ -1512,6 +1526,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) diff --git a/mvcc/index.go b/mvcc/index.go index f0638523f..6caabd60d 100644 --- a/mvcc/index.go +++ b/mvcc/index.go @@ -25,7 +25,7 @@ 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) []revision + 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 @@ -89,7 +89,7 @@ func (ti *treeIndex) keyIndex(keyi *keyIndex) *keyIndex { return nil } -func (ti *treeIndex) visit(key, end []byte, f func(ki *keyIndex)) { +func (ti *treeIndex) visit(key, end []byte, f func(ki *keyIndex) bool) { keyi, endi := &keyIndex{key: key}, &keyIndex{key: end} ti.RLock() @@ -99,25 +99,31 @@ func (ti *treeIndex) visit(key, end []byte, f func(ki *keyIndex)) { if len(endi.key) > 0 && !item.Less(endi) { return false } - f(item.(*keyIndex)) + if !f(item.(*keyIndex)) { + return false + } return true }) } -func (ti *treeIndex) Revisions(key, end []byte, atRev int64) (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) { + ti.visit(key, end, func(ki *keyIndex) bool { if rev, _, _, err := ki.get(ti.lg, atRev); err == nil { - revs = append(revs, rev) + 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) int { @@ -129,10 +135,11 @@ func (ti *treeIndex) CountRevisions(key, end []byte, atRev int64) int { return 1 } total := 0 - ti.visit(key, end, func(ki *keyIndex) { + ti.visit(key, end, func(ki *keyIndex) bool { if _, _, _, err := ki.get(ti.lg, atRev); err == nil { total++ } + return true }) return total } @@ -145,11 +152,12 @@ func (ti *treeIndex) Range(key, end []byte, atRev int64) (keys [][]byte, revs [] } return [][]byte{key}, []revision{rev} } - ti.visit(key, end, func(ki *keyIndex) { + ti.visit(key, end, func(ki *keyIndex) bool { if rev, _, _, err := ki.get(ti.lg, atRev); err == nil { revs = append(revs, rev) keys = append(keys, ki.key) } + return true }) return keys, revs } diff --git a/mvcc/index_test.go b/mvcc/index_test.go index 8befe14cd..ac7a8b182 100644 --- a/mvcc/index_test.go +++ b/mvcc/index_test.go @@ -206,60 +206,80 @@ func TestIndexRevision(t *testing.T) { tests := []struct { key, end []byte atRev int64 + limit int wrevs []revision wcounts int }{ // single key that not found { - []byte("bar"), nil, 6, nil, 0, + []byte("bar"), nil, 6, 0, nil, 0, }, // single key that found { - []byte("foo"), nil, 6, []revision{{main: 6}}, 1, + []byte("foo"), nil, 6, 0, []revision{{main: 6}}, 1, }, - // various range keys, fixed atRev + // various range keys, fixed atRev, unlimited { - []byte("foo"), []byte("foo1"), 6, []revision{{main: 6}}, 1, + []byte("foo"), []byte("foo1"), 6, 0, []revision{{main: 6}}, 1, }, { - []byte("foo"), []byte("foo2"), 6, []revision{{main: 6}, {main: 5}}, 2, + []byte("foo"), []byte("foo2"), 6, 0, []revision{{main: 6}, {main: 5}}, 2, }, { - []byte("foo"), []byte("fop"), 6, []revision{{main: 6}, {main: 5}, {main: 4}}, 3, + []byte("foo"), []byte("fop"), 6, 0, []revision{{main: 6}, {main: 5}, {main: 4}}, 3, }, { - []byte("foo1"), []byte("fop"), 6, []revision{{main: 5}, {main: 4}}, 2, + []byte("foo1"), []byte("fop"), 6, 0, []revision{{main: 5}, {main: 4}}, 2, }, { - []byte("foo2"), []byte("fop"), 6, []revision{{main: 4}}, 1, + []byte("foo2"), []byte("fop"), 6, 0, []revision{{main: 4}}, 1, }, { - []byte("foo3"), []byte("fop"), 6, nil, 0, + []byte("foo3"), []byte("fop"), 6, 0, nil, 0, }, - // fixed range keys, various atRev + // fixed range keys, various atRev, unlimited { - []byte("foo1"), []byte("fop"), 1, nil, 0, + []byte("foo1"), []byte("fop"), 1, 0, nil, 0, }, { - []byte("foo1"), []byte("fop"), 2, []revision{{main: 2}}, 1, + []byte("foo1"), []byte("fop"), 2, 0, []revision{{main: 2}}, 1, }, { - []byte("foo1"), []byte("fop"), 3, []revision{{main: 2}, {main: 3}}, 2, + []byte("foo1"), []byte("fop"), 3, 0, []revision{{main: 2}, {main: 3}}, 2, }, { - []byte("foo1"), []byte("fop"), 4, []revision{{main: 2}, {main: 4}}, 2, + []byte("foo1"), []byte("fop"), 4, 0, []revision{{main: 2}, {main: 4}}, 2, }, { - []byte("foo1"), []byte("fop"), 5, []revision{{main: 5}, {main: 4}}, 2, + []byte("foo1"), []byte("fop"), 5, 0, []revision{{main: 5}, {main: 4}}, 2, }, { - []byte("foo1"), []byte("fop"), 6, []revision{{main: 5}, {main: 4}}, 2, + []byte("foo1"), []byte("fop"), 6, 0, []revision{{main: 5}, {main: 4}}, 2, + }, + // fixed range keys, fixed atRev, various limit + { + []byte("foo"), []byte("fop"), 6, 1, []revision{{main: 6}}, 3, + }, + { + []byte("foo"), []byte("fop"), 6, 2, []revision{{main: 6}, {main: 5}}, 3, + }, + { + []byte("foo"), []byte("fop"), 6, 3, []revision{{main: 6}, {main: 5}, {main: 4}}, 3, + }, + { + []byte("foo"), []byte("fop"), 3, 1, []revision{{main: 1}}, 3, + }, + { + []byte("foo"), []byte("fop"), 3, 2, []revision{{main: 1}, {main: 2}}, 3, + }, + { + []byte("foo"), []byte("fop"), 3, 3, []revision{{main: 1}, {main: 2}, {main: 3}}, 3, }, } for i, tt := range tests { - revs := ti.Revisions(tt.key, tt.end, tt.atRev) + revs, _ := ti.Revisions(tt.key, tt.end, tt.atRev, tt.limit) if !reflect.DeepEqual(revs, tt.wrevs) { - t.Errorf("#%d: revs = %+v, want %+v", i, revs, tt.wrevs) + t.Errorf("#%d limit %d: revs = %+v, want %+v", i, tt.limit, revs, tt.wrevs) } count := ti.CountRevisions(tt.key, tt.end, tt.atRev) if count != tt.wcounts { diff --git a/mvcc/kv_test.go b/mvcc/kv_test.go index 06f82636b..adb71c113 100644 --- a/mvcc/kv_test.go +++ b/mvcc/kv_test.go @@ -219,17 +219,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}) @@ -242,8 +243,12 @@ func testKVRangeLimit(t *testing.T, f rangeFunc) { if r.Rev != wrev { t.Errorf("#%d: rev = %d, want %d", i, r.Rev, wrev) } - if r.Count != len(kvs) { - t.Errorf("#%d: count = %d, want %d", i, r.Count, len(kvs)) + if tt.limit <= 0 || int(tt.limit) > len(kvs) { + if r.Count != len(kvs) { + t.Errorf("#%d: count = %d, want %d", i, r.Count, len(kvs)) + } + } else if r.Count != int(tt.wcounts) { + t.Errorf("#%d: count = %d, want %d", i, r.Count, tt.limit) } } } diff --git a/mvcc/kvstore_test.go b/mvcc/kvstore_test.go index bef9d365b..672878fa6 100644 --- a/mvcc/kvstore_test.go +++ b/mvcc/kvstore_test.go @@ -936,9 +936,12 @@ type fakeIndex struct { indexCompactRespc chan map[revision]struct{} } -func (i *fakeIndex) Revisions(key, end []byte, atRev int64) []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) int { diff --git a/mvcc/kvstore_txn.go b/mvcc/kvstore_txn.go index 5a363b2d1..c7b1f7c54 100644 --- a/mvcc/kvstore_txn.go +++ b/mvcc/kvstore_txn.go @@ -130,10 +130,10 @@ func (tr *storeTxnRead) rangeKeys(key, end []byte, curRev int64, ro RangeOptions 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) + 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(key, end []byte, curRev int64, ro RangeOptions } } 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) {