fix(query): avoid eager UID materialization on posting reads

posting/list.go

@@ -13,6 +13,7 @@ import (
 	"log"
 	"math"
 	"sort"
+	"sync/atomic"
 
 	"github.com/dgryski/go-farm"
 	"github.com/golang/glog"
@@ -65,11 +66,12 @@ const (
 // List stores the in-memory representation of a posting list.
 type List struct {
 	x.SafeMutex
-	key         []byte
-	plist       *pb.PostingList
-	mutationMap *MutableLayer
-	minTs       uint64 // commit timestamp of immutable layer, reject reads before this ts.
-	maxTs       uint64 // max commit timestamp seen for this list.
+	key          []byte
+	plist        *pb.PostingList
+	mutationMap  *MutableLayer
+	minTs        uint64       // commit timestamp of immutable layer, reject reads before this ts.
+	maxTs        uint64       // max commit timestamp seen for this list.
+	uidWarmState atomic.Int32 // One of the uidWarm* states below. See tryStartUidWarm.
 
 	cache []byte
 }
@@ -1053,6 +1055,11 @@ func (l *List) setMutationAfterCommit(startTs, commitTs uint64, pl *pb.PostingLi
 	l.mutationMap.currentUids = nil
 	l.mutationMap.isUidsCalculated = false
 	l.mutationMap.calculatedUids = nil
+	// The list just changed, so a warm that failed against the old state may well succeed against
+	// this one. Reads reach a cached list through this path rather than through a fresh copy, since
+	// remove-on-update defaults to false, so without this a single failure would stick for the life
+	// of the cache entry.
+	l.uidWarmState.CompareAndSwap(uidWarmAbandoned, uidWarmIdle)
 
 	if pl.CommitTs != 0 {
 		l.maxTs = x.Max(l.maxTs, pl.CommitTs)
@@ -1731,37 +1738,98 @@ func (l *List) rollup(readTs uint64, split bool) (*rollupOutput, error) {
 func (l *List) ApproxLen() int {
 	l.RLock()
 	defer l.RUnlock()
+	return l.approxLen()
+}
+
+func (l *List) approxLen() int {
+	l.AssertRLock()
 	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
+// calculateUids materializes the uid slice for l and stores it on l's own mutable layer, reporting
+// whether it did the work.
+//
+// The caller must own l exclusively. This walks the entire list and, for a multi-part list, reads
+// every split from Badger (readListPart), all while holding l's write lock. A list published in the
+// posting-list cache is shared with the commit path, which takes the same write lock from the serial
+// Raft apply loop, so never run this on one; warm a private copy and hand the result over with
+// publishCalculatedUids instead.
+func (l *List) calculateUids() (bool, error) {
+	l.Lock()
+	defer l.Unlock()
+
+	if !l.mutationMap.needsUidWarm() {
+		return false, nil
 	}
-	res := make([]uint64, 0, l.ApproxLen())
 
+	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
+		return false, err
 	}
 
+	l.mutationMap.calculatedUids = res
+	l.mutationMap.isUidsCalculated = true
+	return true, nil
+}
+
+// needsUidWarm reports whether materializing calculatedUids for this layer would be useful. An
+// uncommitted mutation (currentEntries) rules it out because canUseCalculatedUids refuses to serve a
+// slice in that state anyway.
+func (mm *MutableLayer) needsUidWarm() bool {
+	return mm != nil && !mm.isUidsCalculated && mm.currentEntries == nil
+}
+
+// The states of List.uidWarmState. A list starts idle, and copyList builds a fresh List, so a copy
+// never inherits an election or a give-up from the list it was copied from.
+const (
+	uidWarmIdle      int32 = iota // No warm in flight, and none has failed against the list's current state.
+	uidWarmRunning                // A reader is materializing calculatedUids for this list.
+	uidWarmAbandoned              // A warm failed against the current state. Do not retry until it changes.
+)
+
+// tryStartUidWarm elects a single warmer for l, reporting whether the caller won. Materializing
+// calculatedUids is an optimization and never a correctness requirement, so a reader that loses the
+// election serves its read unwarmed rather than queueing behind the winner. Pair a win with
+// finishUidWarm or abandonUidWarm.
+func (l *List) tryStartUidWarm() bool {
+	return l.uidWarmState.CompareAndSwap(uidWarmIdle, uidWarmRunning)
+}
+
+// finishUidWarm releases the election won by tryStartUidWarm, leaving l open to being warmed again.
+// It leaves an abandoned list abandoned, so that it is safe to defer alongside abandonUidWarm.
+func (l *List) finishUidWarm() {
+	l.uidWarmState.CompareAndSwap(uidWarmRunning, uidWarmIdle)
+}
+
+// abandonUidWarm gives up on materializing l's uids, so that no later reader retries it. A walk
+// fails when the list itself cannot be read -- a multi-part list with an unreadable split, say --
+// and retrying costs another full walk plus a Badger read per split, for a slice that reads do not
+// need. setMutationAfterCommit lifts this on the next commit to the list, so a failure suppresses
+// one attempt per state rather than permanently.
+func (l *List) abandonUidWarm() {
+	l.uidWarmState.Store(uidWarmAbandoned)
+}
+
+// publishCalculatedUids installs uids, materialized from a private copy of l taken at
+// committedUidsTime, onto l itself, reporting whether it published. It drops the result if l moved on
+// while the copy was being walked, and it holds l's write lock only for the handover.
+func (l *List) publishCalculatedUids(committedUidsTime uint64, uids []uint64) bool {
 	l.Lock()
 	defer l.Unlock()
 
-	l.mutationMap.calculatedUids = res
-	l.mutationMap.isUidsCalculated = true
+	if !l.mutationMap.needsUidWarm() || l.mutationMap.committedUidsTime != committedUidsTime {
+		return false
+	}
 
-	return nil
+	l.mutationMap.calculatedUids = uids
+	l.mutationMap.isUidsCalculated = true
+	return true
 }
 
 // canUseCalculatedUids reports whether calculatedUids can serve a read at readTs. The slice is
@@ -1784,12 +1852,14 @@ func (l *List) canUseCalculatedUids(readTs uint64) bool {
 // We have to apply the filtering before applying (offset, count).
 // WARNING: Calling this function just to get UIDs is expensive
 func (l *List) Uids(opt ListOptions) (*pb.List, error) {
+	requestedFirst := opt.First
+	bounded := requestedFirst > 0 && requestedFirst < math.MaxInt32
 	if opt.First == 0 {
 		opt.First = math.MaxInt32
 	}
 
 	getUidList := func() (*pb.List, error, bool) {
-		if l.canUseCalculatedUids(opt.ReadTs) {
+		if opt.Intersect == nil && !bounded && l.canUseCalculatedUids(opt.ReadTs) {
 			l.RLock()
 
 			afterIdx := 0
@@ -1811,13 +1881,11 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 			out := &pb.List{Uids: copyArr}
 			l.RUnlock()
 
-			return out, nil, opt.Intersect != nil
+			return out, nil, false
 		}
 		// Pre-assign length to make it faster.
 		l.RLock()
 		defer l.RUnlock()
-		// Use approximate length for initial capacity.
-		res := make([]uint64, 0, l.ApproxLen())
 		out := &pb.List{}
 
 		if l.mutationMap.len() == 0 && opt.Intersect != nil && len(l.plist.Splits) == 0 {
@@ -1828,10 +1896,13 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 			return out, nil, false
 		}
 
+		approxLen := l.approxLen()
+
 		// If we need to intersect and the number of elements are small, in that case it's better to
 		// just check each item is present or not.
-		if opt.Intersect != nil && len(opt.Intersect.Uids) < l.ApproxLen() {
+		if opt.Intersect != nil && len(opt.Intersect.Uids) < approxLen {
 			// Cache the iterator as it makes the search space smaller each time.
+			res := make([]uint64, 0, len(opt.Intersect.Uids))
 			var pitr pIterator
 			for _, uid := range opt.Intersect.Uids {
 				ok, _, err := l.findPostingWithItr(opt.ReadTs, uid, pitr)
@@ -1855,6 +1926,12 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 			uidMax = opt.Intersect.Uids[len(opt.Intersect.Uids)-1]
 		}
 
+		resCap := approxLen
+		if bounded && requestedFirst+1 < resCap {
+			resCap = requestedFirst + 1
+		}
+		res := make([]uint64, 0, resCap)
+
 		err := l.iterate(opt.ReadTs, opt.AfterUid, func(p *pb.Posting) error {
 			if p.PostingType == pb.Posting_REF {
 				if p.Uid < uidMin {
@@ -1865,7 +1942,7 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 				}
 				res = append(res, p.Uid)
 
-				if opt.First != 0 && len(res) > opt.First {
+				if opt.First > 0 && len(res) > opt.First {
 					return ErrStopIteration
 				}
 			}
@@ -1879,24 +1956,25 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 		return out, nil, true
 	}
 
-	// Do The intersection here as it's optimized.
-	out, err, applyIntersectWith := getUidList()
-	if err != nil || !applyIntersectWith || opt.First == 0 {
+	// The bool reports whether the caller still has to intersect what came back and cut it down to
+	// First. A branch that has already intersected, or that needs neither, returns false.
+	out, err, postProcess := getUidList()
+	if err != nil || !postProcess {
 		return out, err
 	}
 
-	if opt.Intersect != nil && applyIntersectWith {
+	if opt.Intersect != nil {
 		algo.IntersectWith(out, opt.Intersect, out)
 	}
 
-	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]
-		}
+	// Only a positive first is applied here, and only because it is paired with the early stop in
+	// the walk above, which makes it an optimization that pays for itself. A negative first has no
+	// early stop, so trimming to the last n buys nothing and is not sound: for an index read the
+	// worker post-filters this list against the real values afterwards -- handleCompareFunction for
+	// a lossy tokenizer, filterGeoFunction for geo -- and rows dropped there cannot come back. The
+	// query layer applies the negative count itself, over the filtered result, with x.PageRange.
+	if opt.First > 0 && len(out.Uids) > opt.First {
+		out.Uids = out.Uids[:opt.First]
 	}
 	return out, nil
 }

posting/list_test.go

@@ -1855,7 +1855,8 @@ func TestCalculatedUidsRespectReadTs(t *testing.T) {
 	require.NoError(t, l.commitMutation(15, 20))
 
 	// Precompute the uid slice. It represents the newest state (commit ts 20).
-	require.NoError(t, l.calculateUids())
+	_, err = l.calculateUids()
+	require.NoError(t, err)
 	require.True(t, l.canUseCalculatedUids(20))
 	require.True(t, l.canUseCalculatedUids(25))
 	require.False(t, l.canUseCalculatedUids(15),
@@ -1874,3 +1875,93 @@ func TestCalculatedUidsRespectReadTs(t *testing.T) {
 	// A read before every commit sees nothing.
 	require.Empty(t, uidsAt(5))
 }
+
+func TestCalculatedUidsSkippedForBoundedReads(t *testing.T) {
+	key := x.DataKey(x.AttrInRootNamespace("calculatedUidsBoundedReads"), 7)
+
+	txn := NewTxn(5)
+	l, err := txn.Get(key)
+	require.NoError(t, err)
+	for _, uid := range []uint64{2, 3, 4} {
+		addMutationHelper(t, l, &pb.DirectedEdge{ValueId: uid}, Set, txn)
+	}
+	require.NoError(t, l.commitMutation(5, 10))
+	_, err = l.calculateUids()
+	require.NoError(t, err)
+	require.True(t, l.canUseCalculatedUids(10))
+
+	l.Lock()
+	l.mutationMap.calculatedUids = []uint64{100, 101}
+	l.Unlock()
+
+	unbounded, err := l.Uids(ListOptions{ReadTs: 10})
+	require.NoError(t, err)
+	require.Equal(t, []uint64{100, 101}, unbounded.Uids)
+
+	workerUnbounded, err := l.Uids(ListOptions{ReadTs: 10, First: math.MaxInt32})
+	require.NoError(t, err)
+	require.Equal(t, []uint64{100, 101}, workerUnbounded.Uids)
+
+	first, err := l.Uids(ListOptions{ReadTs: 10, First: 1})
+	require.NoError(t, err)
+	require.Equal(t, []uint64{2}, first.Uids)
+
+	// A negative first is unbounded as far as the read is concerned, so it is served from the
+	// memoized slice, whole. The poisoned values are what tell the paths apart: {100, 101} could
+	// only have come from the memo, and both of them coming back is what shows it is not trimmed.
+	negative, err := l.Uids(ListOptions{ReadTs: 10, First: -1})
+	require.NoError(t, err)
+	require.Equal(t, []uint64{100, 101}, negative.Uids)
+
+	intersect, err := l.Uids(ListOptions{
+		ReadTs:    10,
+		Intersect: &pb.List{Uids: []uint64{3}},
+	})
+	require.NoError(t, err)
+	require.Equal(t, []uint64{3}, intersect.Uids)
+}
+
+// A negative first must not be applied by Uids at all. For an index read the worker post-filters
+// the list it gets back -- handleCompareFunction against the real values when the tokenizer is
+// lossy, filterGeoFunction against the real geometry -- and a uid trimmed off here is one that
+// filter never sees. The query layer applies the count itself, over the filtered result, with
+// x.PageRange. Both read paths are checked, because they reach the tail differently.
+func TestUidsDoesNotApplyANegativeFirst(t *testing.T) {
+	next := 0
+	build := func(warmed bool) *List {
+		next++
+		key := x.DataKey(x.AttrInRootNamespace(fmt.Sprintf("negativeFirstNotApplied%d", next)), 7)
+		txn := NewTxn(5)
+		l, err := txn.Get(key)
+		require.NoError(t, err)
+		for _, uid := range []uint64{2, 4, 6, 8, 10} {
+			addMutationHelper(t, l, &pb.DirectedEdge{ValueId: uid}, Set, txn)
+		}
+		require.NoError(t, l.commitMutation(5, 10))
+		if warmed {
+			_, err := l.calculateUids()
+			require.NoError(t, err)
+			require.True(t, l.canUseCalculatedUids(10))
+		} else {
+			require.False(t, l.mutationMap.isUidsCalculated)
+		}
+		return l
+	}
+
+	all := []uint64{2, 4, 6, 8, 10}
+	for _, first := range []int{-1, -2, -4, -5, -9, math.MinInt32, math.MinInt} {
+		for _, warmed := range []bool{false, true} {
+			got, err := build(warmed).Uids(ListOptions{ReadTs: 10, First: first})
+			require.NoError(t, err)
+			require.Equal(t, all, got.Uids, "first: %d, warmed: %t", first, warmed)
+		}
+	}
+
+	// AfterUid still applies: it bounds which uids exist for this read, rather than how many of
+	// them to return.
+	for _, warmed := range []bool{false, true} {
+		got, err := build(warmed).Uids(ListOptions{ReadTs: 10, First: -2, AfterUid: 4})
+		require.NoError(t, err)
+		require.Equal(t, []uint64{6, 8, 10}, got.Uids, "warmed: %t", warmed)
+	}
+}