diff --git a/mvcc/index.go b/mvcc/index.go index c9b0d1831..a88ef46e8 100644 --- a/mvcc/index.go +++ b/mvcc/index.go @@ -26,6 +26,7 @@ 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 + 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 @@ -119,6 +120,23 @@ func (ti *treeIndex) Revisions(key, end []byte, atRev int64) (revs []revision) { return revs } +func (ti *treeIndex) CountRevisions(key, end []byte, atRev int64) int { + if end == nil { + _, _, _, err := ti.Get(key, atRev) + if err != nil { + return 0 + } + return 1 + } + total := 0 + ti.visit(key, end, func(ki *keyIndex) { + if _, _, _, err := ki.get(ti.lg, atRev); err == nil { + total++ + } + }) + return total +} + func (ti *treeIndex) Range(key, end []byte, atRev int64) (keys [][]byte, revs []revision) { if end == nil { rev, _, _, err := ti.Get(key, atRev) diff --git a/mvcc/kvstore_test.go b/mvcc/kvstore_test.go index 6d2e1aa20..ff858fc7c 100644 --- a/mvcc/kvstore_test.go +++ b/mvcc/kvstore_test.go @@ -941,6 +941,11 @@ func (i *fakeIndex) Revisions(key, end []byte, atRev int64) []revision { return rev } +func (i *fakeIndex) CountRevisions(key, end []byte, atRev int64) int { + _, rev := i.Range(key, end, atRev) + return len(rev) +} + func (i *fakeIndex) Get(key []byte, atRev int64) (rev, created revision, ver int64, err error) { i.Recorder.Record(testutil.Action{Name: "get", Params: []interface{}{key, atRev}}) r := <-i.indexGetRespc diff --git a/mvcc/kvstore_txn.go b/mvcc/kvstore_txn.go index e89ddbee4..2b89dd50b 100644 --- a/mvcc/kvstore_txn.go +++ b/mvcc/kvstore_txn.go @@ -125,15 +125,16 @@ func (tr *storeTxnRead) rangeKeys(key, end []byte, curRev int64, ro RangeOptions if rev < tr.s.compactMainRev { return &RangeResult{KVs: nil, Count: -1, Rev: 0}, ErrCompacted } - + if ro.Count { + 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) tr.trace.Step("range keys from in-memory index tree") if len(revpairs) == 0 { return &RangeResult{KVs: nil, Count: 0, Rev: curRev}, nil } - if ro.Count { - return &RangeResult{KVs: nil, Count: len(revpairs), Rev: curRev}, nil - } limit := int(ro.Limit) if limit <= 0 || limit > len(revpairs) {