Fix a bug where deletions weren't updating the raw Kind index.

impl/memory/datastore_query_execution.go

 // Copyright 2015 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 package memory
 
 import (
         "bytes"
         "errors"
         "fmt"
 
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/gae/service/datastore/serialize"
         "github.com/luci/luci-go/common/data/cmpbin"
         "github.com/luci/luci-go/common/data/stringset"
 )
 
+// errQueryEntrySkipped is a sentinel error returned by queryStrategy handlers
+// if they decided not to process the entry in their handle().
+var errQueryEntrySkipped = errors.New("query result skipped")
+
 type queryStrategy interface {
         // handle applies the strategy to the embedded user callback.
         //   - rawData is the slice of encoded Properties from the index row
         //     (correctly de-inverted).
         //   - decodedProps is the slice of decoded Properties from the index row
         //   - key is the decoded Key from the index row (the last item in rawData and
         //     decodedProps)
         //   - gc is the getCursor function to be passed to the user's callback
-        handle(rawData [][]byte, decodedProps []ds.Property, key *ds.Key, gc func() (ds.Cursor, error)) error
+        handle(rawData [][]byte, decodedProps []ds.Property, key *ds.Key,
+                cb ds.RawRunCB, gc func() (ds.Cursor, error)) error
 }
 
 type projectionLookup struct {
         suffixIndex  int
         propertyName string
 }
 
 type projectionStrategy struct {
-        cb ds.RawRunCB
-
         project  []projectionLookup
         distinct stringset.Set
 }
 
-func newProjectionStrategy(fq *ds.FinalizedQuery, rq *reducedQuery, cb ds.RawRunCB) queryStrategy {
+func newProjectionStrategy(fq *ds.FinalizedQuery, rq *reducedQuery) queryStrategy {
         proj := fq.Project()
 
         projectionLookups := make([]projectionLookup, len(proj))
         for i, prop := range proj {
                 projectionLookups[i].propertyName = prop
                 lookupErr := fmt.Errorf("planning a strategy for an unfulfillable query?")
                 for j, col := range rq.suffixFormat {
                         if col.Property == prop {
                                 projectionLookups[i].suffixIndex = j
                                 lookupErr = nil
                                 break
                         }
                 }
                 impossible(lookupErr)
         }
-        ret := &projectionStrategy{cb: cb, project: projectionLookups}
+        ret := &projectionStrategy{project: projectionLookups}
         if fq.Distinct() {
                 ret.distinct = stringset.New(0)
         }
         return ret
 }
 
-func (s *projectionStrategy) handle(rawData [][]byte, decodedProps []ds.Property, key *ds.Key, gc func() (ds.Cursor, error)) error {
+func (s *projectionStrategy) handle(rawData [][]byte, decodedProps []ds.Property, key *ds.Key,
+        cb ds.RawRunCB, gc func() (ds.Cursor, error)) error {
+
         projectedRaw := [][]byte(nil)
         if s.distinct != nil {
                 projectedRaw = make([][]byte, len(decodedProps))
         }
         pmap := make(ds.PropertyMap, len(s.project))
         for i, p := range s.project {
                 if s.distinct != nil {
                         projectedRaw[i] = rawData[p.suffixIndex]
                 }
                 pmap[p.propertyName] = decodedProps[p.suffixIndex]
         }
         if s.distinct != nil {
                 if !s.distinct.Add(string(serialize.Join(projectedRaw...))) {
                         return nil
                 }
         }
-        return s.cb(key, pmap, gc)
+        return cb(key, pmap, gc)
 }
 
 type keysOnlyStrategy struct {
-        cb ds.RawRunCB
-
         dedup stringset.Set
 }
 
-func (s *keysOnlyStrategy) handle(rawData [][]byte, _ []ds.Property, key *ds.Key, gc func() (ds.Cursor, error)) error {
+func (s *keysOnlyStrategy) handle(rawData [][]byte, _ []ds.Property, key *ds.Key,
+        cb ds.RawRunCB, gc func() (ds.Cursor, error)) error {
+
         if !s.dedup.Add(string(rawData[len(rawData)-1])) {
                 return nil
         }
-        return s.cb(key, nil, gc)
+        return cb(key, nil, gc)
 }
 
 type normalStrategy struct {
-        cb ds.RawRunCB
-
-        kc    ds.KeyContext
-        head  memCollection
-        dedup stringset.Set
+        kc      ds.KeyContext
+        isCount bool
+        head    memCollection
+        dedup   stringset.Set
 }
 
-func newNormalStrategy(kc ds.KeyContext, cb ds.RawRunCB, head memStore) queryStrategy {
+func newNormalStrategy(kc ds.KeyContext, isCount bool, head memStore) queryStrategy {
         coll := head.GetCollection("ents:" + kc.Namespace)
         if coll == nil {
                 return nil
         }
-        return &normalStrategy{cb, kc, coll, stringset.New(0)}
+        return &normalStrategy{kc, isCount, coll, stringset.New(0)}
 }
 
-func (s *normalStrategy) handle(rawData [][]byte, _ []ds.Property, key *ds.Key, gc func() (ds.Cursor, error)) error {
+func (s *normalStrategy) handle(rawData [][]byte, _ []ds.Property, key *ds.Key,
+        cb ds.RawRunCB, gc func() (ds.Cursor, error)) error {
+
         rawKey := rawData[len(rawData)-1]
         if !s.dedup.Add(string(rawKey)) {
                 return nil
         }
 
         rawEnt := s.head.Get(rawKey)
         if rawEnt == nil {
                 // entity doesn't exist at head
-                return nil
+                return errQueryEntrySkipped

[dnj, 2016/11/11 08:37:53]

This is the source of the problem. If the callback is skipped b/c entity doesn't
exist at HEAD, "limit" and "offset" are still advanced, counting against it.

         }
-        pm, err := serialize.ReadPropertyMap(bytes.NewBuffer(rawEnt), serialize.WithoutContext, s.kc)
-        memoryCorruption(err)
 
-        return s.cb(key, pm, gc)
+        var pm ds.PropertyMap
+        if !s.isCount {
+                var err error
+                if pm, err = serialize.ReadPropertyMap(bytes.NewBuffer(rawEnt), serialize.WithoutContext, s.kc); err != nil {
+                        memoryCorruption(err)
+                }
+        }
+
+        return cb(key, pm, gc)
 }
 
-func pickQueryStrategy(fq *ds.FinalizedQuery, rq *reducedQuery, cb ds.RawRunCB, head memStore) queryStrategy {
+func pickQueryStrategy(fq *ds.FinalizedQuery, rq *reducedQuery, isCount bool, head memStore) queryStrategy {
         if fq.KeysOnly() {
-                return &keysOnlyStrategy{cb, stringset.New(0)}
+                return &keysOnlyStrategy{stringset.New(0)}
         }
         if len(fq.Project()) > 0 {
-                return newProjectionStrategy(fq, rq, cb)
+                return newProjectionStrategy(fq, rq)
         }
-        return newNormalStrategy(rq.kc, cb, head)
+        return newNormalStrategy(rq.kc, isCount, head)
 }
 
 func parseSuffix(aid, ns string, suffixFormat []ds.IndexColumn, suffix []byte, count int) (raw [][]byte, decoded []ds.Property) {
         buf := serialize.Invertible(bytes.NewBuffer(suffix))
         decoded = make([]ds.Property, len(suffixFormat))
         raw = make([][]byte, len(suffixFormat))
 
         err := error(nil)
         kc := ds.MkKeyContext(aid, ns)
         for i := range decoded {
@@ skipping 11 matching lines @@
                 suffix = suffix[offset:]
                 if needInvert {
                         raw[i] = serialize.Invert(raw[i])
                 }
         }
 
         return
 }
 
 func countQuery(fq *ds.FinalizedQuery, kc ds.KeyContext, isTxn bool, idx, head memStore) (ret int64, err error) {
-        if len(fq.Project()) == 0 && !fq.KeysOnly() {

[dnj, 2016/11/11 08:37:53]

So maybe I'm reading things wrong, but I think Count should behave the same as
Get. Making it KeysOnly here allows it to return keys from stale indexes that
wouldn't be returned by Get, making it incorrect.

-                fq, err = fq.Original().KeysOnly(true).Finalize()
-                if err != nil {
-                        return
-                }
-        }
-        err = executeQuery(fq, kc, isTxn, idx, head, func(_ *ds.Key, _ ds.PropertyMap, _ ds.CursorCB) error {
+        err = executeQuery(fq, kc, isTxn, true, idx, head, func(_ *ds.Key, _ ds.PropertyMap, _ ds.CursorCB) error {
                 ret++
                 return nil
         })
         return
 }
 
 func executeNamespaceQuery(fq *ds.FinalizedQuery, kc ds.KeyContext, head memStore, cb ds.RawRunCB) error {
         // these objects have no properties, so any filters on properties cause an
         // empty result.
         if len(fq.EqFilters()) > 0 || len(fq.Project()) > 0 || len(fq.Orders()) > 1 {
@@ skipping 32 matching lines @@
                 } else {
                         k = kc.MakeKey("__namespace__", ns)
                 }
                 if err := cb(k, nil, cursFn); err != nil {
                         return err
                 }
         }
         return nil
 }
 
-func executeQuery(fq *ds.FinalizedQuery, kc ds.KeyContext, isTxn bool, idx, head memStore, cb ds.RawRunCB) error {
+func executeQuery(fq *ds.FinalizedQuery, kc ds.KeyContext, isTxn, isCount bool, idx, head memStore, cb ds.RawRunCB) error {
         rq, err := reduce(fq, kc, isTxn)
         if err == ds.ErrNullQuery {
                 return nil
         }
         if err != nil {
                 return err
         }
 
         if rq.kind == "__namespace__" {
                 return executeNamespaceQuery(fq, kc, head, cb)
         }
 
         idxs, err := getIndexes(rq, idx)
         if err == ds.ErrNullQuery {
                 return nil
         }
         if err != nil {
                 return err
         }
 
-        strategy := pickQueryStrategy(fq, rq, cb, head)
+        strategy := pickQueryStrategy(fq, rq, isCount, head)
         if strategy == nil {
                 // e.g. the normalStrategy found that there were NO entities in the current
                 // namespace.
                 return nil
         }
 
         offset, _ := fq.Offset()
         limit, hasLimit := fq.Limit()
 
         cursorPrefix := []byte(nil)
@@ skipping 10 matching lines @@
                                 }
                                 cursorPrefix = buf.Bytes()
                         }
                         // TODO(riannucci): Do we need to decrement suffix instead of increment
                         // if we're sorting by __key__ DESCENDING?
                         return queryCursor(serialize.Join(cursorPrefix, increment(suffix))), nil
                 }
         }
 
         return multiIterate(idxs, func(suffix []byte) error {
+                iterCB := cb
+
+                // Apply our limit/offset.
                 if offset > 0 {
-                        offset--
-                        return nil
+                        // Don't actually callback for this entry.
+                        iterCB = func(*ds.Key, ds.PropertyMap, ds.CursorCB) error { return nil }
                 }
                 if hasLimit {
                         if limit <= 0 {
                                 return ds.Stop
                         }
-                        limit--
                 }
 
                 rawData, decodedProps := parseSuffix(kc.AppID, kc.Namespace, rq.suffixFormat, suffix, -1)
 
                 keyProp := decodedProps[len(decodedProps)-1]
                 if keyProp.Type() != ds.PTKey {
                         impossible(fmt.Errorf("decoded index row doesn't end with a Key: %#v", keyProp))
                 }
 
-                return strategy.handle(
+                err := strategy.handle(
                         rawData, decodedProps, keyProp.Value().(*ds.Key),
-                        getCursorFn(suffix))
+                        iterCB, getCursorFn(suffix))
+
+                if err == errQueryEntrySkipped {
+                        return nil
+                }
+
+                // We processed this entry; advance our limit/offset.
+                if offset > 0 {
+                        offset--
+                } else if hasLimit {
+                        limit--
+                }
+                return err
         })
 }