Add Consistent(bool) to Testing interface

impl/memory/datastore_data.go

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package memory
 
 import (
         "bytes"
         "fmt"
         "sync"
         "sync/atomic"
 
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/gae/service/datastore/serialize"
         "github.com/luci/luci-go/common/errors"
         "golang.org/x/net/context"
 )
 
 //////////////////////////////// dataStoreData /////////////////////////////////
 
 type dataStoreData struct {
         rwlock sync.RWMutex
         // See README.md for head schema.
         head *memStore
         snap *memStore
         // For testing, see SetTransactionRetryCount.
         txnFakeRetry int
+        // true means that head always == snap
+        consistent bool
 }
 
 var (
         _ = memContextObj((*dataStoreData)(nil))
         _ = sync.Locker((*dataStoreData)(nil))
 )
 
 func newDataStoreData() *dataStoreData {
         head := newMemStore()
         return &dataStoreData{
                 head: head,
                 snap: head.Snapshot(), // empty but better than a nil pointer.
         }
 }
 
 func (d *dataStoreData) Lock() {
         d.rwlock.Lock()
 }
 
 func (d *dataStoreData) Unlock() {
         d.rwlock.Unlock()
 }
 
+func (d *dataStoreData) setTxnRetry(count int) {
+        d.Lock()

[Vadim Sh., 2015/09/24 18:40:11]

oops.. didn't know it's under lock

[iannucci, 2015/09/24 18:59:30]

is ok :)

+        defer d.Unlock()
+        d.txnFakeRetry = count
+}
+
+func (d *dataStoreData) setConsistent(always bool) {
+        d.Lock()
+        defer d.Unlock()
+
+        d.consistent = always
+        if d.consistent {
+                d.snap = d.head.Snapshot()
+        }
+}
+
 func (d *dataStoreData) getQuerySnaps(consistent bool) (idx, head *memStore) {
         d.rwlock.RLock()
         defer d.rwlock.RUnlock()
+        if d.consistent {
+                // snap is already a consistent snapshot of head
+                return d.snap, d.snap
+        }
+
         head = d.head.Snapshot()
         if consistent {
                 idx = head
         } else {
                 idx = d.snap
         }
         return
 }
 
 func (d *dataStoreData) takeSnapshot() *memStore {
         d.rwlock.RLock()
         defer d.rwlock.RUnlock()
+        if d.consistent {
+                return d.snap
+        }
         return d.head.Snapshot()
 }
 
 func (d *dataStoreData) setSnapshot(snap *memStore) {
         d.rwlock.Lock()
         defer d.rwlock.Unlock()
+        if d.consistent {
+                return
+        }
         d.snap = snap
 }
 
 func (d *dataStoreData) catchupIndexes() {
         d.rwlock.Lock()
         defer d.rwlock.Unlock()
+        if d.consistent {
+                return
+        }
         d.snap = d.head.Snapshot()
 }
 
 /////////////////////////// indexes(dataStoreData) ////////////////////////////
 
 func groupMetaKey(key *ds.Key) []byte {
         return keyBytes(ds.NewKey("", "", "__entity_group__", "", 1, key.Root()))
 }
 
 func groupIDsKey(key *ds.Key) []byte {
@@ skipping 64 matching lines @@
 
 func (d *dataStoreData) fixKeyLocked(ents *memCollection, key *ds.Key) *ds.Key {
         if key.Incomplete() {
                 id := d.allocateIDsLocked(ents, key, 1)
                 key = ds.NewKey(key.AppID(), key.Namespace(), key.Kind(), "", id, key.Parent())
         }
         return key
 }
 
 func (d *dataStoreData) putMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.PutMultiCB) {
-        ents := d.mutableEnts(keys[0].Namespace())
+        ns := keys[0].Namespace()
+        ents := d.mutableEnts(ns)
 
         for i, k := range keys {
                 pmap, _ := vals[i].Save(false)
                 dataBytes := serialize.ToBytes(pmap)
 
                 k, err := func() (ret *ds.Key, err error) {
                         d.Lock()
                         defer d.Unlock()
 
                         ret = d.fixKeyLocked(ents, k)
                         incrementLocked(ents, groupMetaKey(ret), 1)
 
                         old := ents.Get(keyBytes(ret))
                         oldPM := ds.PropertyMap(nil)
                         if old != nil {
-                                if oldPM, err = rpmWoCtx(old, ret.Namespace()); err != nil {
+                                if oldPM, err = rpmWoCtx(old, ns); err != nil {
                                         return
                                 }
                         }
                         updateIndexes(d.head, ret, oldPM, pmap)
                         ents.Set(keyBytes(ret), dataBytes)
+                        if d.consistent {
+                                d.snap = d.head.Snapshot()
+                        }
                         return
                 }()
                 if cb != nil {
                         cb(k, err)
                 }
         }
 }
 
 func getMultiInner(keys []*ds.Key, cb ds.GetMultiCB, getColl func() (*memCollection, error)) error {
         ents, err := getColl()
@@ skipping 20 matching lines @@
 
 func (d *dataStoreData) getMulti(keys []*ds.Key, cb ds.GetMultiCB) error {
         return getMultiInner(keys, cb, func() (*memCollection, error) {
                 s := d.takeSnapshot()
 
                 return s.GetCollection("ents:" + keys[0].Namespace()), nil
         })
 }
 
 func (d *dataStoreData) delMulti(keys []*ds.Key, cb ds.DeleteMultiCB) {
-        toDel := make([][]byte, 0, len(keys))
-        for _, k := range keys {
-                toDel = append(toDel, keyBytes(k))
-        }
         ns := keys[0].Namespace()
+        ents := d.mutableEnts(ns)
 
-        d.rwlock.Lock()
-        defer d.rwlock.Unlock()
+        if ents != nil {
+                for _, k := range keys {
+                        err := func() error {
+                                kb := keyBytes(k)
 
-        ents := d.head.GetCollection("ents:" + ns)
+                                d.Lock()
+                                defer d.Unlock()
 
-        for i, k := range keys {
-                if ents != nil {
-                        incrementLocked(ents, groupMetaKey(k), 1)
-                        kb := toDel[i]
-                        if old := ents.Get(kb); old != nil {
-                                oldPM, err := rpmWoCtx(old, ns)
-                                if err != nil {
-                                        if cb != nil {
-                                                cb(err)
+                                incrementLocked(ents, groupMetaKey(k), 1)
+                                if old := ents.Get(kb); old != nil {
+                                        oldPM, err := rpmWoCtx(old, ns)
+                                        if err != nil {
+                                                return err
                                         }
-                                        continue
+                                        updateIndexes(d.head, k, oldPM, nil)
+                                        ents.Delete(kb)
+                                        if d.consistent {
+                                                d.snap = d.head.Snapshot()
+                                        }
                                 }
-                                updateIndexes(d.head, k, oldPM, nil)
-                                ents.Delete(kb)
+                                return nil
+                        }()
+                        if cb != nil {

[Vadim Sh., 2015/09/24 18:40:11]

I think you need to call cb() len(keys) number of times even if 'ents == nil'.

[iannucci, 2015/09/24 18:59:30]

oh, good catch.

+                                cb(err)
                         }
                 }
-                if cb != nil {
-                        cb(nil)
-                }
         }
 }
 
 func (d *dataStoreData) canApplyTxn(obj memContextObj) bool {
         // TODO(riannucci): implement with Flush/FlushRevert for persistance.
 
         txn := obj.(*txnDataStoreData)
         for rk, muts := range txn.muts {
                 if len(muts) == 0 { // read-only
                         continue
@@ skipping 186 matching lines @@
 
 func rpmWoCtx(data []byte, ns string) (ds.PropertyMap, error) {
         return serialize.ReadPropertyMap(bytes.NewBuffer(data),
                 serialize.WithoutContext, globalAppID, ns)
 }
 
 func rpm(data []byte) (ds.PropertyMap, error) {
         return serialize.ReadPropertyMap(bytes.NewBuffer(data),
                 serialize.WithContext, "", "")
 }