diff --git a/posting/list.go b/posting/list.go index d96ac5247d5..6df8e9cf4f4 100644 --- a/posting/list.go +++ b/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 } diff --git a/posting/list_test.go b/posting/list_test.go index a27e9416402..07cfaf1718c 100644 --- a/posting/list_test.go +++ b/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) + } +} diff --git a/posting/mvcc.go b/posting/mvcc.go index badf13b9fa3..aaea3f55626 100644 --- a/posting/mvcc.go +++ b/posting/mvcc.go @@ -409,6 +409,10 @@ func (ml *MemoryLayer) UpdateMaxCost(maxCost int64) { ml.cache.data.UpdateMaxCost(maxCost) } +func (ml *MemoryLayer) hasCache() bool { + return ml != nil && ml.cache != nil && ml.cache.data != nil +} + type IterateDiskArgs struct { Prefix []byte Prefetch bool @@ -761,19 +765,90 @@ func (c *CachePL) Set(l *List, readTs uint64) { } } -func (ml *MemoryLayer) readFromCache(key []byte, readTs uint64) *List { +func (ml *MemoryLayer) readFromCache(key []byte, readTs uint64, readUids bool) *List { cacheItem, ok := ml.cache.get(key) + if !ok || cacheItem.list == nil { + return nil + } + cached := cacheItem.list + cached.RLock() // Issue #9597 fix: Cache is only valid if minTs <= readTs AND maxTs >= readTs. // If maxTs < readTs, the cache is missing mutations committed after maxTs. - if ok && cacheItem.list != nil && cacheItem.list.minTs <= readTs && cacheItem.list.maxTs >= readTs { - cacheItem.list.RLock() - lCopy := copyList(cacheItem.list) - cacheItem.list.RUnlock() - checkForRollup(key, lCopy) - return lCopy + if cached.minTs > readTs || cached.maxTs < readTs { + cached.RUnlock() + return nil } - return nil + needsWarm := readUids && cached.mutationMap.needsUidWarm() + lCopy := copyList(cached) + cached.RUnlock() + + if needsWarm { + ml.warmCachedUids(key, cacheItem, lCopy) + } + checkForRollup(key, lCopy) + return lCopy +} + +// warmCachedUids materializes the uid slice for a cached posting list and publishes it back to the +// cached entry, so that later readers of the same key are served from it instead of walking the list +// again. +// +// The work happens on lCopy, which this reader owns, rather than on the published list. A warm walks +// the whole list and reads split parts from Badger, and the commit path takes the published list's +// write lock from the serial Raft apply loop (commitOrAbort in worker/draft.go, ahead of the +// ProcessDelta that releases waiting reads), so holding that lock across the walk would stall every +// commit for the group. Exactly one reader warms at a time; the rest serve their read unwarmed and +// recompute as they did before. +// +// Iterating lCopy is safe even though its mutable layer shares maps with the published list, but +// only because updateItemInCache, the one caller that commits into a cached list, passes +// refresh=true, and that rebuilds committedEntries and committedUids before writing. A refresh=false +// commit on a published list would write into the maps a copy like this one is reading, which Go +// reports as a fatal concurrent map access rather than a race the detector might miss. +func (ml *MemoryLayer) warmCachedUids(key []byte, cacheItem *CachePL, lCopy *List) { + cached := cacheItem.list + if !cached.tryStartUidWarm() { + return + } + defer cached.finishUidWarm() + + calculated, err := lCopy.calculateUids() + if err != nil { + // Warming is an optimization, so it must not be the thing that fails the read. Often the + // read fails anyway, on the same unreadable split, once it walks the list itself -- but a + // transient error need not recur, and a read bounded by First or AfterUid may never reach + // the part that could not be read. + // + // Give up on this list rather than letting the next reader repeat the walk. Everyone who + // gets here holds a cache hit on a list that cannot be read, so otherwise both the wasted + // walk and this log line recur at read QPS. Abandoning caps them at one per commit to the + // key, which is the bound doRollup gets from its own per-key dedupe. + cached.abandonUidWarm() + glog.Warningf("Giving up on materializing calculated UIDs for key [%x]: %v", key, err) + return + } + if !calculated { + return + } + if !cached.publishCalculatedUids(lCopy.mutationMap.committedUidsTime, + lCopy.mutationMap.calculatedUids) { + return + } + // Re-set the entry so ristretto accounts for the calculated UID slice. + ml.resetIfCurrent(key, cacheItem) +} + +// resetIfCurrent re-sets cacheItem under key so that its ristretto cost is recomputed, but only if it +// is still the live entry for that key. A plain set would resurrect an entry that was evicted or +// dropped by a rollup while the caller was working, and the caller may have been working for as long +// as a warm takes. +func (ml *MemoryLayer) resetIfCurrent(key []byte, cacheItem *CachePL) bool { + if current, ok := ml.cache.get(key); !ok || current != cacheItem { + return false + } + ml.cache.set(key, cacheItem) + return true } func (ml *MemoryLayer) readFromDisk(key []byte, pstore *badger.DB, readTs uint64, readUids bool) (*List, error) { @@ -792,9 +867,13 @@ func (ml *MemoryLayer) readFromDisk(key []byte, pstore *badger.DB, readTs uint64 if err != nil { return l, err } - if readUids { - if err := l.calculateUids(); err != nil { - return nil, err + if readUids && ml.hasCache() { + if _, err := l.calculateUids(); err != nil { + // Same as on the cache path: the list itself read fine, so serve it unmaterialized rather + // than failing the read for the sake of an optimization. This one keeps its warning + // unconditionally, because it runs per cache miss rather than per read, and it is where + // an operator should first see that a list has stopped being readable. + glog.Warningf("Error materializing calculated UIDs for key [%x]: %v", key, err) } } return l, nil @@ -814,7 +893,7 @@ func (ml *MemoryLayer) ReadData(key []byte, pstore *badger.DB, readTs uint64, re // 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. - l := ml.readFromCache(key, readTs) + l := ml.readFromCache(key, readTs, readUids) if l != nil { l.mutationMap.setTs(readTs) return l, nil diff --git a/posting/mvcc_test.go b/posting/mvcc_test.go index 3a0553c886c..1fab24d9520 100644 --- a/posting/mvcc_test.go +++ b/posting/mvcc_test.go @@ -18,6 +18,7 @@ import ( "google.golang.org/protobuf/proto" "github.com/dgraph-io/badger/v4" + "github.com/dgraph-io/dgraph/v25/codec" "github.com/dgraph-io/dgraph/v25/protos/pb" "github.com/dgraph-io/dgraph/v25/schema" "github.com/dgraph-io/dgraph/v25/x" @@ -273,6 +274,335 @@ func TestCacheStaleWhenMaxTsLessThanReadTs(t *testing.T) { require.True(t, hasUid2, "UID 2 missing - cache returned stale data (maxTs < readTs)") } +func TestReadUidsHonorsPostingListCache(t *testing.T) { + require.NoError(t, pstore.DropAll()) + useMemoryLayer(t, 0) + + attr := x.AttrInRootNamespace("readUidsCache") + key := x.DataKey(attr, 1) + addEdgeToUID(t, attr, 1, 2, 1, 2) + addEdgeToUID(t, attr, 1, 3, 3, 4) + + l, err := getNew(key, pstore, math.MaxUint64, true) + require.NoError(t, err) + require.False(t, l.mutationMap.isUidsCalculated, + "readUids should not materialize UIDs when posting-list cache is disabled") + + useMemoryLayer(t, 10<<20) + l, err = getNew(key, pstore, math.MaxUint64, true) + require.NoError(t, err) + require.True(t, l.mutationMap.isUidsCalculated, + "readUids should still warm calculated UIDs when posting-list cache is enabled") +} + +func TestReadUidsWarmsCachedPostingList(t *testing.T) { + require.NoError(t, pstore.DropAll()) + useMemoryLayer(t, 10<<20) + + attr := x.AttrInRootNamespace("readUidsCacheHit") + key := x.DataKey(attr, 1) + addEdgeToUID(t, attr, 1, 2, 1, 2) + addEdgeToUID(t, attr, 1, 3, 3, 4) + + l, err := getNew(key, pstore, math.MaxUint64, false) + require.NoError(t, err) + require.False(t, l.mutationMap.isUidsCalculated) + MemLayerInstance.wait() + + cacheItem, ok := MemLayerInstance.cache.get(key) + require.True(t, ok) + require.False(t, cacheItem.list.mutationMap.isUidsCalculated) + readTs := cacheItem.list.maxTs + cacheSizeBefore := cacheItem.list.ApproximateSize() + remainingCostBefore := MemLayerInstance.cache.data.RemainingCost() + + l, err = getNew(key, pstore, readTs, true) + require.NoError(t, err) + require.True(t, l.mutationMap.isUidsCalculated) + MemLayerInstance.wait() + + cacheItem, ok = MemLayerInstance.cache.get(key) + require.True(t, ok) + require.True(t, cacheItem.list.mutationMap.isUidsCalculated) + cacheSizeAfter := cacheItem.list.ApproximateSize() + require.Greater(t, cacheSizeAfter, cacheSizeBefore) + require.Equal(t, int64(cacheSizeAfter-cacheSizeBefore), + remainingCostBefore-MemLayerInstance.cache.data.RemainingCost()) +} + +// cacheUidList publishes a posting list of uidCount refs under key and returns its cached entry. +func cacheUidList(t *testing.T, key []byte, uidCount uint64) *CachePL { + encoder := codec.Encoder{BlockSize: 256} + for uid := uint64(1); uid <= uidCount; uid++ { + encoder.Add(uid) + } + pack := encoder.Done() + t.Cleanup(func() { + codec.FreePack(pack) + }) + + MemLayerInstance.saveInCache(key, &List{ + key: key, + plist: &pb.PostingList{Pack: pack}, + mutationMap: newMutableLayer(), + minTs: 1, + maxTs: 1, + }) + MemLayerInstance.wait() + + cacheItem, ok := MemLayerInstance.cache.get(key) + require.True(t, ok) + return cacheItem +} + +// useMemoryLayer swaps in a posting-list cache of the given size for the duration of the test. +func useMemoryLayer(t *testing.T, cacheSize int64) { + orig := MemLayerInstance + MemLayerInstance = initMemoryLayer(cacheSize, false) + t.Cleanup(func() { + MemLayerInstance = orig + }) +} + +func TestReadUidsWarmsCachedPostingListConcurrently(t *testing.T) { + require.NoError(t, pstore.DropAll()) + useMemoryLayer(t, 10<<20) + + const uidCount = 4096 + key := x.DataKey(x.AttrInRootNamespace("readUidsConcurrentCacheHit"), 1) + cacheUidList(t, key, uidCount) + + const readers = 32 + start := make(chan struct{}) + type readResult struct { + list *List + err error + } + results := make(chan readResult, readers) + for range readers { + go func() { + <-start + l, err := getNew(key, pstore, 1, true) + results <- readResult{list: l, err: err} + }() + } + close(start) + + // Drain every reader before asserting, so that a failure cannot run the cleanup that restores + // MemLayerInstance while readers are still using it. + collected := make([]readResult, 0, readers) + for range readers { + collected = append(collected, <-results) + } + + // Only one reader materializes the uids; the rest are served an unwarmed copy. Every one of them + // has to return the same answer either way. + for _, result := range collected { + require.NoError(t, result.err) + uids, err := result.list.Uids(ListOptions{ReadTs: 1}) + require.NoError(t, err) + require.Len(t, uids.Uids, uidCount) + } + + MemLayerInstance.wait() + cacheItem, ok := MemLayerInstance.cache.get(key) + require.True(t, ok) + require.True(t, cacheItem.list.canUseCalculatedUids(1)) + require.Len(t, cacheItem.list.mutationMap.calculatedUids, uidCount) +} + +func TestWarmCachedUidsWalksWithoutTheCachedListsWriteLock(t *testing.T) { + require.NoError(t, pstore.DropAll()) + useMemoryLayer(t, 10<<20) + + const uidCount = 4096 + key := x.DataKey(x.AttrInRootNamespace("readUidsWarmWithoutWriteLock"), 1) + cacheItem := cacheUidList(t, key, uidCount) + cached := cacheItem.list + + cached.RLock() + lCopy := copyList(cached) + cached.RUnlock() + + warmed := make(chan struct{}) + go func() { + defer close(warmed) + MemLayerInstance.warmCachedUids(key, cacheItem, lCopy) + }() + + func() { + // Hold the published list's read lock for the whole walk. A warm that took its write lock + // would deadlock here, and would also shut out every other reader, because Go's RWMutex stops + // admitting readers once a writer is queued. The commit path takes that same write lock from + // the serial Raft apply loop, so the walk has to stay off it. + cached.RLock() + defer cached.RUnlock() + + require.Eventually(t, func() bool { + lCopy.RLock() + defer lCopy.RUnlock() + return lCopy.mutationMap.isUidsCalculated + }, 30*time.Second, time.Millisecond, + "materializing calculated UIDs blocked on the published list's write lock") + }() + + <-warmed + + // The handover is the only part that needs the write lock, and it hands over the same slice. + // require.Same rather than require.Equal: Equal follows both pointers and compares the uids + // they address, so it passes for two distinct arrays that happen to start with the same uid. + require.True(t, cached.canUseCalculatedUids(1)) + require.Len(t, cached.mutationMap.calculatedUids, uidCount) + require.Same(t, &lCopy.mutationMap.calculatedUids[0], &cached.mutationMap.calculatedUids[0]) +} + +func TestReadUidsDoesNotWaitForAnInFlightWarm(t *testing.T) { + require.NoError(t, pstore.DropAll()) + useMemoryLayer(t, 10<<20) + + const uidCount = 512 + key := x.DataKey(x.AttrInRootNamespace("readUidsInFlightWarm"), 1) + cacheItem := cacheUidList(t, key, uidCount) + + // Stand in for another reader's warm that is still walking the list. + require.True(t, cacheItem.list.tryStartUidWarm()) + t.Cleanup(cacheItem.list.finishUidWarm) + + l, err := getNew(key, pstore, 1, true) + require.NoError(t, err) + require.False(t, l.canUseCalculatedUids(1), + "a reader that loses the warm election should be served an unwarmed copy instead of waiting") + + // Unwarmed is still correct: the reader walks the list as it did before the optimization. + uids, err := l.Uids(ListOptions{ReadTs: 1}) + require.NoError(t, err) + require.Len(t, uids.Uids, uidCount) +} + +func TestWarmCachedUidsGivesUpAfterAFailedWalk(t *testing.T) { + require.NoError(t, pstore.DropAll()) + useMemoryLayer(t, 10<<20) + + key := x.DataKey(x.AttrInRootNamespace("readUidsWarmFailure"), 1) + + // A list that claims a split which was never written, so every walk of it fails on the Badger + // read for that part. This is the shape of the failure the give-up exists for. + MemLayerInstance.saveInCache(key, &List{ + key: key, + plist: &pb.PostingList{Splits: []uint64{1}}, + mutationMap: newMutableLayer(), + minTs: 1, + maxTs: 1, + }) + MemLayerInstance.wait() + + cacheItem, ok := MemLayerInstance.cache.get(key) + require.True(t, ok) + cached := cacheItem.list + + // The read still succeeds, unmaterialized, and the failure is not reported to the caller. + l, err := getNew(key, pstore, 1, true) + require.NoError(t, err) + require.False(t, l.canUseCalculatedUids(1)) + require.Equal(t, uidWarmAbandoned, cached.uidWarmState.Load(), + "a failed walk should give up on the entry instead of leaving it open to retries") + + // A later reader neither retries the walk nor wins an election, so the cost is paid once for + // the entry rather than once per read. + require.False(t, cached.tryStartUidWarm()) + l, err = getNew(key, pstore, 1, true) + require.NoError(t, err) + require.False(t, l.canUseCalculatedUids(1)) + require.Equal(t, uidWarmAbandoned, cached.uidWarmState.Load()) + + // Giving up is per List, so a fresh entry for the same key starts willing to try again. + replacement := cacheUidList(t, key, 8) + require.Equal(t, uidWarmIdle, replacement.list.uidWarmState.Load()) +} + +func TestCommitLetsAnAbandonedWarmTryAgain(t *testing.T) { + l := &List{key: x.DataKey(x.AttrInRootNamespace("abandonedWarmReset"), 1), mutationMap: newMutableLayer()} + + require.True(t, l.tryStartUidWarm()) + l.abandonUidWarm() + require.False(t, l.tryStartUidWarm()) + + // A commit applies in place on the cached list, because remove-on-update defaults to false, so + // the give-up has to lift here or a single failure would stick for the life of the entry. + l.setMutationAfterCommit(5, 10, &pb.PostingList{ + Postings: []*pb.Posting{{Uid: 2, PostingType: pb.Posting_REF, Op: Set}}, + }, true) + + require.Equal(t, uidWarmIdle, l.uidWarmState.Load(), + "a commit changes the list, so a warm that failed against the old state should get another try") + require.True(t, l.tryStartUidWarm()) +} + +func TestFinishUidWarmLeavesAnAbandonedListAlone(t *testing.T) { + l := &List{mutationMap: newMutableLayer()} + + require.True(t, l.tryStartUidWarm()) + l.abandonUidWarm() + // finishUidWarm is deferred on the same path that abandons, so it must not reopen the list. + l.finishUidWarm() + require.Equal(t, uidWarmAbandoned, l.uidWarmState.Load()) + require.False(t, l.tryStartUidWarm()) + + other := &List{mutationMap: newMutableLayer()} + require.True(t, other.tryStartUidWarm()) + other.finishUidWarm() + require.Equal(t, uidWarmIdle, other.uidWarmState.Load()) + require.True(t, other.tryStartUidWarm()) +} + +func TestPublishCalculatedUidsDropsAStaleWalk(t *testing.T) { + l := &List{mutationMap: newMutableLayer()} + l.mutationMap.committedUidsTime = 7 + + // A commit landed while the copy was being walked, so the result describes an older snapshot. + require.False(t, l.publishCalculatedUids(5, []uint64{1, 2})) + require.False(t, l.mutationMap.isUidsCalculated) + require.Empty(t, l.mutationMap.calculatedUids) + + require.True(t, l.publishCalculatedUids(7, []uint64{1, 2})) + require.True(t, l.mutationMap.isUidsCalculated) + require.Equal(t, []uint64{1, 2}, l.mutationMap.calculatedUids) + + // An already warmed list is left alone. + require.False(t, l.publishCalculatedUids(7, []uint64{3})) + require.Equal(t, []uint64{1, 2}, l.mutationMap.calculatedUids) +} + +func TestResetIfCurrentLeavesADroppedEntryOut(t *testing.T) { + require.NoError(t, pstore.DropAll()) + useMemoryLayer(t, 10<<20) + + key := x.DataKey(x.AttrInRootNamespace("resetIfCurrent"), 1) + cacheItem := cacheUidList(t, key, 8) + + require.True(t, MemLayerInstance.resetIfCurrent(key, cacheItem)) + MemLayerInstance.wait() + _, ok := MemLayerInstance.cache.get(key) + require.True(t, ok) + + // A rollup drops the entry while the warm is still running. Re-setting it would resurrect it. + MemLayerInstance.del(key) + MemLayerInstance.wait() + require.False(t, MemLayerInstance.resetIfCurrent(key, cacheItem)) + MemLayerInstance.wait() + _, ok = MemLayerInstance.cache.get(key) + require.False(t, ok, "a dropped cache entry should stay dropped") + + // Same for an entry that was replaced by a newer read of the key. + replacement := cacheUidList(t, key, 16) + require.NotSame(t, cacheItem, replacement) + require.False(t, MemLayerInstance.resetIfCurrent(key, cacheItem)) + MemLayerInstance.wait() + current, ok := MemLayerInstance.cache.get(key) + require.True(t, ok) + require.Same(t, replacement, current) +} + func TestPostingListRead(t *testing.T) { attr := x.AttrInRootNamespace("emptypl") key := x.DataKey(attr, 1) diff --git a/posting/size.go b/posting/size.go index 5de90a0b43c..78c6764f717 100644 --- a/posting/size.go +++ b/posting/size.go @@ -28,6 +28,7 @@ func (l *List) ApproximateSize() uint64 { 1*8 + // plist pointer consists of 1 word. 1*8 + // mutation map pointer consists of 1 word. 2*8 + // minTs and maxTs take 1 word each. + 1*8 + // uidWarmState takes 1 word including alignment. 3*8 + // array take 3 words. so key array is 3 words. 3*8 + // array take 3 words. so cache array is 3 words. 1*8 // So far 11 words, in order to round the slab we're adding one more word. diff --git a/query/pagination_test.go b/query/pagination_test.go new file mode 100644 index 00000000000..55638a8bbaa --- /dev/null +++ b/query/pagination_test.go @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package query + +import ( + "math" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/dgraph-io/dgraph/v25/protos/pb" +) + +// A pushed-down first is applied by posting.(*List).Uids to each posting list the worker reads, so +// it is only sound for a function that reads one list per result. A function that reads a list per +// token and then intersects them needs the whole of each list, because the first (or last) n of an +// intersection is not the intersection of the first (or last) n. calculatePaginationParams is what +// keeps those functions off the pushdown, and this pins that. +func TestPaginationPushdownExcludesIntersectingFunctions(t *testing.T) { + // Reading a list per token and intersecting: a pushdown would drop matches. + for _, fn := range []string{"regexp", "alloftext", "allofterms", "match", "ngram"} { + t.Run(fn, func(t *testing.T) { + for _, count := range []int{5, -5} { + sg := &SubGraph{ + SrcFunc: &Function{Name: fn}, + Params: params{Count: count, Offset: 7}, + } + first, offset := calculatePaginationParams(sg) + require.Equal(t, int32(math.MaxInt32), first, "count: %d", count) + require.Zero(t, offset, "count: %d", count) + } + }) + } + + // Reading one list, or reading several and merging them: a pushdown is sound, because the + // first (or last) n of a union is contained in the union of each list's first (or last) n, + // and the query layer takes its own slice of the merged result afterwards. + for _, fn := range []string{"eq", "anyofterms", "uid_in"} { + t.Run(fn, func(t *testing.T) { + sg := &SubGraph{ + SrcFunc: &Function{Name: fn}, + Params: params{Count: 5, Offset: 7}, + } + first, offset := calculatePaginationParams(sg) + require.Equal(t, int32(5), first) + require.Equal(t, int32(7), offset) + }) + } +} + +// A filter, an order, or no count at all all mean the whole list is needed, whatever the function. +func TestPaginationPushdownNeedsTheWholeListSometimes(t *testing.T) { + unbounded := func(sg *SubGraph) { + first, offset := calculatePaginationParams(sg) + require.Equal(t, int32(math.MaxInt32), first) + require.Zero(t, offset) + } + + unbounded(&SubGraph{Params: params{Count: 0, Offset: 7}}) + unbounded(&SubGraph{Params: params{Count: 5, Offset: 7}, Filters: []*SubGraph{{}}}) + unbounded(&SubGraph{Params: params{Count: 5, Offset: 7, Order: []*pb.Order{{Attr: "name"}}}}) +} diff --git a/worker/precalculate_uids_test.go b/worker/precalculate_uids_test.go new file mode 100644 index 00000000000..ed8d37903ac --- /dev/null +++ b/worker/precalculate_uids_test.go @@ -0,0 +1,119 @@ +/* + * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package worker + +import ( + "math" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/dgraph-io/dgraph/v25/posting" + "github.com/dgraph-io/dgraph/v25/protos/pb" +) + +func TestShouldPrecalculateUids(t *testing.T) { + tests := []struct { + name string + q *pb.Query + srcFn *functionContext + facetsTree *facetsTree + opts posting.ListOptions + want bool + }{ + { + name: "allows unbounded uid reads", + q: &pb.Query{}, + srcFn: &functionContext{fnType: standardFn}, + want: true, + }, + { + name: "allows internal unbounded first sentinel", + q: &pb.Query{}, + srcFn: &functionContext{fnType: standardFn}, + opts: posting.ListOptions{First: math.MaxInt32}, + want: true, + }, + { + name: "skips bounded first reads", + q: &pb.Query{}, + srcFn: &functionContext{fnType: standardFn}, + opts: posting.ListOptions{First: 10}, + }, + { + name: "skips intersect reads", + q: &pb.Query{}, + srcFn: &functionContext{fnType: standardFn}, + opts: posting.ListOptions{Intersect: &pb.List{Uids: []uint64{1}}}, + }, + { + name: "skips count reads", + q: &pb.Query{DoCount: true}, + srcFn: &functionContext{fnType: standardFn}, + }, + { + name: "skips facet reads", + q: &pb.Query{FacetParam: &pb.FacetParams{}}, + srcFn: &functionContext{fnType: standardFn}, + }, + { + name: "skips facet filter reads", + q: &pb.Query{}, + srcFn: &functionContext{fnType: standardFn}, + facetsTree: &facetsTree{}, + }, + { + name: "skips compare scalar reads", + q: &pb.Query{}, + srcFn: &functionContext{fnType: compareScalarFn}, + }, + { + name: "skips has function reads", + q: &pb.Query{}, + srcFn: &functionContext{fnType: hasFn}, + }, + { + name: "skips uid in reads", + q: &pb.Query{}, + srcFn: &functionContext{fnType: uidInFn}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.want, shouldPrecalculateUids(tc.q, tc.srcFn, tc.facetsTree, tc.opts)) + }) + } +} + +func TestUidReadFirst(t *testing.T) { + tests := []struct { + name string + first int32 + offset int32 + want int + }{ + {name: "bounded read includes offset", first: 10, offset: 5, want: 15}, + {name: "unbounded read keeps zero", first: 0, offset: 5, want: 0}, + {name: "negative first ignores offset", first: -2, offset: 5, want: -2}, + {name: "unbounded sentinel ignores offset", first: math.MaxInt32, offset: 5, + want: math.MaxInt32}, + {name: "bounded read near the sentinel does not overflow int32", + first: math.MaxInt32 - 1, offset: 5, want: math.MaxInt32 + 4}, + // x.PageRange clamps a negative offset to 0, so the pushdown has to as well. Adding it + // would read fewer uids than were asked for, and a large enough one would go negative, + // which means "take from the back" by the time it reaches List.Uids. + {name: "small negative offset does not shrink the read", first: 10, offset: -1, want: 10}, + {name: "large negative offset does not invert the read", first: 10, offset: -100, want: 10}, + {name: "negative first still ignores a negative offset", first: -2, offset: -100, want: -2}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.want, uidReadFirst(&pb.Query{First: tc.first, Offset: tc.offset})) + }) + } +} diff --git a/worker/task.go b/worker/task.go index a59bff64fcc..f142f43cfff 100644 --- a/worker/task.go +++ b/worker/task.go @@ -778,6 +778,34 @@ func retrieveUidsAndFacets(args funcArgs, pl *posting.List, facetsTree *facetsTr return uidList, fcsList, nil } +func shouldPrecalculateUids(q *pb.Query, srcFn *functionContext, facetsTree *facetsTree, + opts posting.ListOptions) bool { + if q.DoCount || q.FacetParam != nil || facetsTree != nil { + return false + } + if opts.Intersect != nil || (opts.First > 0 && opts.First < math.MaxInt32) { + return false + } + switch srcFn.fnType { + case compareScalarFn, hasFn, uidInFn: + return false + default: + return true + } +} + +// Negative first ignores offset, while zero and MaxInt32 both mean that the worker must read +// the whole list. Only positive bounded reads can safely stop after first + offset entries. +func uidReadFirst(q *pb.Query) int { + if q.First <= 0 || q.First == math.MaxInt32 { + return int(q.First) + } + // A negative offset counts as no offset, which is how x.PageRange reads one when it paginates + // the result. Adding it instead would read fewer uids than were asked for, and a large enough + // one would push the read negative, which drops the bound altogether. + return int(int64(q.First) + int64(max(q.Offset, 0))) +} + // This function handles operations on uid posting lists. Index keys, reverse keys and some data // keys store uid posting lists. func (qs *queryState) handleUidPostings( @@ -824,6 +852,7 @@ func (qs *queryState) handleUidPostings( isList := schema.State().IsList(q.Attr) outputs := make([]*pb.Result, numGo) + precalculateUids := shouldPrecalculateUids(q, srcFn, facetsTree, opts) eg, egCtx := errgroup.WithContext(ctx) calculate := func(start, end int) error { @@ -855,7 +884,13 @@ func (qs *queryState) handleUidPostings( } // Get or create the posting list for an entity, attribute combination. - pl, err := qs.cache.GetUids(key) + var pl *posting.List + var err error + if precalculateUids { + pl, err = qs.cache.GetUids(key) + } else { + pl, err = qs.cache.Get(key) + } if err != nil { return err } @@ -936,7 +971,7 @@ func (qs *queryState) handleUidPostings( ReadTs: args.q.ReadTs, AfterUid: 0, Intersect: reqList, - First: int(args.q.First + args.q.Offset), + First: uidReadFirst(args.q), } plist, err := pl.Uids(topts) if err != nil { @@ -1120,7 +1155,7 @@ func (qs *queryState) helpProcessTask(ctx context.Context, q *pb.Query, gid uint opts := posting.ListOptions{ ReadTs: q.ReadTs, AfterUid: q.AfterUid, - First: int(q.First + q.Offset), + First: uidReadFirst(q), } // If we have srcFunc and Uids, it means its a filter. So we intersect. if srcFn.fnType != notAFunction && q.UidList != nil && len(q.UidList.Uids) > 0 {