perf(core): Cache uids array in ristretto (#9430)

Author: harshil-goel

dgraph/cmd/debug/run.go

@@ -603,7 +603,7 @@ func printKeys(db *badger.DB) {
 		}
 
 		var sz, deltaCount int64
-		pl, err := posting.GetNew(key, db, opt.readTs)
+		pl, err := posting.GetNew(key, db, opt.readTs, false)
 		if err == nil {
 			pl.RLock()
 			c := pl.GetLength(math.MaxUint64)

dgraph/cmd/live/batch.go

@@ -450,11 +450,14 @@ loop:
 		}
 
 		drain()
-		time.Sleep(100 * time.Millisecond)
 	}
+
+	for len(buffer) != 0 {
+		drain()
+	}
+
 	fmt.Printf("Looped %d times over buffered requests.\n", loops)
 
-	drain()
 }
 
 func (l *loader) printCounters() {

posting/index_test.go

@@ -172,7 +172,7 @@ func TestTokensTable(t *testing.T) {
 
 	attr := x.AttrInRootNamespace("name")
 	key := x.DataKey(attr, 1)
-	l, err := getNew(key, ps, math.MaxUint64)
+	l, err := getNew(key, ps, math.MaxUint64, false)
 	require.NoError(t, err)
 
 	edge := &pb.DirectedEdge{

posting/list.go

@@ -100,6 +100,10 @@ type MutableLayer struct {
 
 	// We also cache some things required for us to update currentEntries faster
 	currentUids map[uint64]int // Stores the uid to index mapping in the currentEntries posting list
+
+	// Cache for calculated UIDS
+	isUidsCalculated bool
+	calculatedUids   []uint64
 }
 
 func newMutableLayer() *MutableLayer {
@@ -110,6 +114,8 @@ func newMutableLayer() *MutableLayer {
 		length:            math.MaxInt,
 		committedUids:     make(map[uint64]*pb.Posting),
 		committedUidsTime: math.MaxUint64,
+		isUidsCalculated:  false,
+		calculatedUids:    []uint64{},
 	}
 }
 
@@ -136,6 +142,8 @@ func (mm *MutableLayer) clone() *MutableLayer {
 		length:            mm.length,
 		lastEntry:         mm.lastEntry,
 		committedUidsTime: mm.committedUidsTime,
+		isUidsCalculated:  mm.isUidsCalculated,
+		calculatedUids:    mm.calculatedUids,
 	}
 }
 
@@ -153,6 +161,8 @@ func (mm *MutableLayer) setCurrentEntries(ts uint64, pl *pb.PostingList) {
 	mm.readTs = ts
 	mm.currentEntries = pl
 	clear(mm.currentUids)
+	mm.isUidsCalculated = false
+	mm.calculatedUids = nil
 	mm.deleteAllMarker = math.MaxUint64
 	mm.populateUidMap(pl)
 }
@@ -770,6 +780,9 @@ func (l *List) updateMutationLayer(mpost *pb.Posting, singleUidUpdate, hasCountI
 		l.mutationMap.currentEntries = &pb.PostingList{}
 	}
 
+	l.mutationMap.isUidsCalculated = false
+	l.mutationMap.calculatedUids = nil
+
 	if singleUidUpdate {
 		// This handles the special case when adding a value to predicates of type uid.
 		// The current value should be deleted in favor of this value. This needs to
@@ -1010,6 +1023,8 @@ func (l *List) setMutationAfterCommit(startTs, commitTs uint64, pl *pb.PostingLi
 	l.mutationMap.currentEntries = nil
 	l.mutationMap.readTs = 0
 	l.mutationMap.currentUids = nil
+	l.mutationMap.isUidsCalculated = false
+	l.mutationMap.calculatedUids = nil
 
 	if pl.CommitTs != 0 {
 		l.maxTs = x.Max(l.maxTs, pl.CommitTs)
@@ -1691,6 +1706,43 @@ func (l *List) ApproxLen() int {
 	return l.mutationMap.len() + codec.ApproxLen(l.plist.Pack)
 }
 
+func (l *List) calculateUids() error {
+	l.RLock()
+	if l.mutationMap == nil || l.mutationMap.isUidsCalculated {
+		l.RUnlock()
+		return nil
+	}
+	res := make([]uint64, 0, l.ApproxLen())
+
+	err := l.iterate(l.mutationMap.committedUidsTime, 0, func(p *pb.Posting) error {
+		if p.PostingType == pb.Posting_REF {
+			res = append(res, p.Uid)
+		}
+		return nil
+	})
+
+	l.RUnlock()
+
+	if err != nil {
+		return err
+	}
+
+	l.Lock()
+	defer l.Unlock()
+
+	l.mutationMap.calculatedUids = res
+	l.mutationMap.isUidsCalculated = true
+
+	return nil
+}
+
+func (l *List) canUseCalculatedUids() bool {
+	if l.mutationMap == nil {
+		return false
+	}
+	return l.mutationMap.isUidsCalculated && l.mutationMap.currentEntries == nil
+}
+
 // Uids returns the UIDs given some query params.
 // We have to apply the filtering before applying (offset, count).
 // WARNING: Calling this function just to get UIDs is expensive
@@ -1700,6 +1752,30 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 	}
 
 	getUidList := func() (*pb.List, error, bool) {
+		if l.canUseCalculatedUids() {
+			l.RLock()
+
+			afterIdx := 0
+
+			if opt.AfterUid != 0 {
+				after := sort.Search(len(l.mutationMap.calculatedUids), func(i int) bool {
+					return l.mutationMap.calculatedUids[i] > opt.AfterUid
+				})
+				if after >= len(l.mutationMap.calculatedUids) {
+					l.RUnlock()
+					return &pb.List{}, nil, false
+				}
+
+				afterIdx = after
+			}
+
+			copyArr := make([]uint64, len(l.mutationMap.calculatedUids)-afterIdx)
+			copy(copyArr, l.mutationMap.calculatedUids[afterIdx:])
+			out := &pb.List{Uids: copyArr}
+			l.RUnlock()
+
+			return out, nil, opt.Intersect != nil
+		}
 		// Pre-assign length to make it faster.
 		l.RLock()
 		defer l.RUnlock()
@@ -1752,13 +1828,7 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 				}
 				res = append(res, p.Uid)
 
-				if opt.First < 0 {
-					// We need the last N.
-					// TODO: This could be optimized by only considering some of the last UidBlocks.
-					if len(res) > -opt.First {
-						res = res[1:]
-					}
-				} else if len(res) > opt.First {
+				if opt.First != 0 && len(res) > opt.First {
 					return ErrStopIteration
 				}
 			}
@@ -1774,19 +1844,22 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 
 	// Do The intersection here as it's optimized.
 	out, err, applyIntersectWith := getUidList()
-	if err != nil || !applyIntersectWith {
+	if err != nil || !applyIntersectWith || opt.First == 0 {
 		return out, err
 	}
 
-	lenBefore := len(out.Uids)
-	if opt.Intersect != nil {
+	if opt.Intersect != nil && applyIntersectWith {
 		algo.IntersectWith(out, opt.Intersect, out)
 	}
-	lenAfter := len(out.Uids)
-	if lenBefore-lenAfter > 0 {
-		// If we see this log, that means that iterate is going over too many elements that it doesn't need to
-		glog.V(3).Infof("Retrieved a list. length before intersection: %d, length after: %d, extra"+
-			" elements: %d", lenBefore, lenAfter, lenBefore-lenAfter)
+
+	if opt.First != 0 {
+		if opt.First < 0 {
+			if len(out.Uids) > -opt.First {
+				out.Uids = out.Uids[(len(out.Uids) + opt.First):]
+			}
+		} else if len(out.Uids) > opt.First {
+			out.Uids = out.Uids[:opt.First]
+		}
 	}
 	return out, nil
 }

posting/list_test.go

@@ -500,7 +500,7 @@ func TestReadSingleValue(t *testing.T) {
 	// and GetSingeValueForKey works without an issue.
 
 	key := x.DataKey(x.AttrInRootNamespace("value"), 1240)
-	ol, err := getNew(key, ps, math.MaxUint64)
+	ol, err := getNew(key, ps, math.MaxUint64, false)
 	require.NoError(t, err)
 	N := uint64(10000)
 	for i := uint64(2); i <= N; i += 2 {
@@ -525,7 +525,7 @@ func TestReadSingleValue(t *testing.T) {
 			require.NoError(t, writePostingListToDisk(kvs))
 			// Delete item from global cache before reading, as we are not updating the cache in the test
 			memoryLayer.del(key)
-			ol, err = getNew(key, ps, math.MaxUint64)
+			ol, err = getNew(key, ps, math.MaxUint64, false)
 			require.NoError(t, err)
 		}
 

posting/lists.go

@@ -57,7 +57,7 @@ func Cleanup() {
 // GetNoStore returns the list stored in the key or creates a new one if it doesn't exist.
 // It does not store the list in any cache.
 func GetNoStore(key []byte, readTs uint64) (rlist *List, err error) {
-	return getNew(key, pstore, readTs)
+	return getNew(key, pstore, readTs, false)
 }
 
 // LocalCache stores a cache of posting lists and deltas.
@@ -258,13 +258,13 @@ func (lc *LocalCache) SetIfAbsent(key string, updated *List) *List {
 	return updated
 }
 
-func (lc *LocalCache) getInternal(key []byte, readFromDisk bool) (*List, error) {
+func (lc *LocalCache) getInternal(key []byte, readFromDisk, readUids bool) (*List, error) {
 	skey := string(key)
 	getNewPlistNil := func() (*List, error) {
 		lc.RLock()
 		defer lc.RUnlock()
 		if lc.plists == nil {
-			return getNew(key, pstore, lc.startTs)
+			return getNew(key, pstore, lc.startTs, readUids)
 		}
 		if l, ok := lc.plists[skey]; ok {
 			return l, nil
@@ -279,7 +279,7 @@ func (lc *LocalCache) getInternal(key []byte, readFromDisk bool) (*List, error)
 	var pl *List
 	if readFromDisk {
 		var err error
-		pl, err = getNew(key, pstore, lc.startTs)
+		pl, err = getNew(key, pstore, lc.startTs, readUids)
 		if err != nil {
 			return nil, err
 		}
@@ -390,14 +390,18 @@ func (lc *LocalCache) GetSinglePosting(key []byte) (*pb.PostingList, error) {
 
 // Get retrieves the cached version of the list associated with the given key.
 func (lc *LocalCache) Get(key []byte) (*List, error) {
-	return lc.getInternal(key, true)
+	return lc.getInternal(key, true, false)
+}
+
+func (lc *LocalCache) GetUids(key []byte) (*List, error) {
+	return lc.getInternal(key, true, true)
 }
 
 // GetFromDelta gets the cached version of the list without reading from disk
 // and only applies the existing deltas. This is used in situations where the
 // posting list will only be modified and not read (e.g adding index mutations).
 func (lc *LocalCache) GetFromDelta(key []byte) (*List, error) {
-	return lc.getInternal(key, false)
+	return lc.getInternal(key, false, false)
 }
 
 // UpdateDeltasAndDiscardLists updates the delta cache before removing the stored posting lists.

posting/lmap_test.go

@@ -18,7 +18,7 @@ func BenchmarkGet(b *testing.B) {
 		for pb.Next() {
 			// i := uint64(rand.Int63())
 			_ = uint64(rand.Int63())
-			_, _ = getNew(key, nil, math.MaxUint64)
+			_, _ = getNew(key, nil, math.MaxUint64, false)
 			// lmap.Get(i)
 		}
 	})
@@ -30,7 +30,7 @@ func BenchmarkGetLinear(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		k := uint64(i)
 		if _, ok := m[k]; !ok {
-			l, err := getNew(key, nil, math.MaxUint64)
+			l, err := getNew(key, nil, math.MaxUint64, false)
 			if err != nil {
 				b.Error(err)
 			}

posting/mvcc.go

@@ -636,6 +636,7 @@ func ReadPostingList(key []byte, it *badger.Iterator) (*List, error) {
 		}
 		it.Next()
 	}
+
 	return l, nil
 }
 
@@ -671,7 +672,7 @@ func (ml *MemoryLayer) readFromCache(key []byte, readTs uint64) *List {
 	return nil
 }
 
-func (ml *MemoryLayer) readFromDisk(key []byte, pstore *badger.DB, readTs uint64) (*List, error) {
+func (ml *MemoryLayer) readFromDisk(key []byte, pstore *badger.DB, readTs uint64, readUids bool) (*List, error) {
 	txn := pstore.NewTransactionAt(readTs, false)
 	defer txn.Discard()
 
@@ -687,6 +688,14 @@ func (ml *MemoryLayer) readFromDisk(key []byte, pstore *badger.DB, readTs uint64
 	if err != nil {
 		return l, err
 	}
+	if readUids {
+		if l.mutationMap == nil {
+			l.mutationMap = newMutableLayer()
+		}
+		if err := l.calculateUids(); err != nil {
+			return nil, err
+		}
+	}
 	return l, nil
 }
 
@@ -700,7 +709,7 @@ func (ml *MemoryLayer) saveInCache(key []byte, l *List) {
 	ml.cache.set(key, cacheItem)
 }
 
-func (ml *MemoryLayer) ReadData(key []byte, pstore *badger.DB, readTs uint64) (*List, error) {
+func (ml *MemoryLayer) ReadData(key []byte, pstore *badger.DB, readTs uint64, readUids bool) (*List, error) {
 	// We first try to read the data from cache, if it is present. If it's not present, then we would read the
 	// latest data from the disk. This would get stored in the cache. If this read has a minTs > readTs then
 	// we would have to read the correct timestamp from the disk.
@@ -709,7 +718,7 @@ func (ml *MemoryLayer) ReadData(key []byte, pstore *badger.DB, readTs uint64) (*
 		l.mutationMap.setTs(readTs)
 		return l, nil
 	}
-	l, err := ml.readFromDisk(key, pstore, math.MaxUint64)
+	l, err := ml.readFromDisk(key, pstore, math.MaxUint64, readUids)
 	if err != nil {
 		return nil, err
 	}
@@ -719,7 +728,7 @@ func (ml *MemoryLayer) ReadData(key []byte, pstore *badger.DB, readTs uint64) (*
 		return l, nil
 	}
 
-	l, err = ml.readFromDisk(key, pstore, readTs)
+	l, err = ml.readFromDisk(key, pstore, readTs, readUids)
 	if err != nil {
 		return nil, err
 	}
@@ -728,16 +737,16 @@ func (ml *MemoryLayer) ReadData(key []byte, pstore *badger.DB, readTs uint64) (*
 	return l, nil
 }
 
-func GetNew(key []byte, pstore *badger.DB, readTs uint64) (*List, error) {
-	return getNew(key, pstore, readTs)
+func GetNew(key []byte, pstore *badger.DB, readTs uint64, readUids bool) (*List, error) {
+	return getNew(key, pstore, readTs, readUids)
 }
 
-func getNew(key []byte, pstore *badger.DB, readTs uint64) (*List, error) {
+func getNew(key []byte, pstore *badger.DB, readTs uint64, readUids bool) (*List, error) {
 	if pstore.IsClosed() {
 		return nil, badger.ErrDBClosed
 	}
 
-	l, err := memoryLayer.ReadData(key, pstore, readTs)
+	l, err := memoryLayer.ReadData(key, pstore, readTs, readUids)
 	if err != nil {
 		return l, err
 	}