Add the ability to disable special automatic entities in impl/memory

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"
@@ skipping 12 matching lines @@
         // See README.md for head schema.
         head *memStore
         snap *memStore
         // For testing, see SetTransactionRetryCount.
         txnFakeRetry int
         // true means that head always == snap
         consistent bool
         // true means that queries with insufficent indexes will pause to add them
         // and then continue instead of failing.
         autoIndex bool
+        // true means that all of the __...__ keys which are normally automatically
+        // maintained will be omitted. This also means that Put with an incomplete
+        // key will become an error.
+        disableSpecialEntities bool
 }
 
 var (
         _ = memContextObj((*dataStoreData)(nil))
         _ = sync.Locker((*dataStoreData)(nil))
 )
 
 func newDataStoreData() *dataStoreData {
         head := newMemStore()
         return &dataStoreData{
@@ skipping 52 matching lines @@
         d.rwlock.RUnlock()
 
         if !ai {
                 return false
         }
 
         d.addIndexes(mi.ns, []*ds.IndexDefinition{mi.Missing})
         return true
 }
 
+func (d *dataStoreData) setDisableSpecialEntities(enabled bool) {
+        d.Lock()
+        defer d.Unlock()
+        d.disableSpecialEntities = true
+}
+
+func (d *dataStoreData) getDisableSpecialEntities() bool {
+        d.rwlock.RLock()
+        defer d.rwlock.RUnlock()
+        return d.disableSpecialEntities
+}
+
 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 {
@@ skipping 78 matching lines @@
         defer d.Unlock()
 
         coll := "ents:" + ns
         ents := d.head.GetCollection(coll)
         if ents == nil {
                 ents = d.head.SetCollection(coll, nil)
         }
         return ents
 }
 
-func (d *dataStoreData) allocateIDs(incomplete *ds.Key, n int) int64 {
+func (d *dataStoreData) allocateIDs(incomplete *ds.Key, n int) (int64, error) {
         ents := d.mutableEnts(incomplete.Namespace())
 
         d.Lock()
         defer d.Unlock()
         return d.allocateIDsLocked(ents, incomplete, n)
 }
 
-func (d *dataStoreData) allocateIDsLocked(ents *memCollection, incomplete *ds.Key, n int) int64 {
+func (d *dataStoreData) allocateIDsLocked(ents *memCollection, incomplete *ds.Key, n int) (int64, error) {
+        if d.disableSpecialEntities {
+                return 0, errors.New("disableSpecialEntities is true so allocateIDs is disabled")
+        }
+
         idKey := []byte(nil)
         if incomplete.Parent() == nil {
                 idKey = rootIDsKey(incomplete.Kind())
         } else {
                 idKey = groupIDsKey(incomplete)
         }
-        return incrementLocked(ents, idKey, n)
+        return incrementLocked(ents, idKey, n), nil
 }
 
-func (d *dataStoreData) fixKeyLocked(ents *memCollection, key *ds.Key) *ds.Key {
+func (d *dataStoreData) fixKeyLocked(ents *memCollection, key *ds.Key) (*ds.Key, error) {
         if key.Incomplete() {
-                id := d.allocateIDsLocked(ents, key, 1)
+                id, err := d.allocateIDsLocked(ents, key, 1)
+                if err != nil {
+                        return key, err
+                }
                 key = ds.NewKey(key.AppID(), key.Namespace(), key.Kind(), "", id, key.Parent())
         }
-        return key
+        return key, nil
 }
 
 func (d *dataStoreData) putMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.PutMultiCB) {
         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)
+                        ret, err = d.fixKeyLocked(ents, k)
+                        if err != nil {
+                                return
+                        }
+                        if !d.disableSpecialEntities {
+                                incrementLocked(ents, groupMetaKey(ret), 1)
+                        }
 
                         old := ents.Get(keyBytes(ret))
                         oldPM := ds.PropertyMap(nil)
                         if old != nil {
                                 if oldPM, err = rpmWoCtx(old, ns); err != nil {
                                         return
                                 }
                         }
                         updateIndexes(d.head, ret, oldPM, pmap)
                         ents.Set(keyBytes(ret), dataBytes)
@@ skipping 44 matching lines @@
         ents := d.mutableEnts(ns)
 
         if ents != nil {
                 for _, k := range keys {
                         err := func() error {
                                 kb := keyBytes(k)
 
                                 d.Lock()
                                 defer d.Unlock()
 
-                                incrementLocked(ents, groupMetaKey(k), 1)
+                                if !d.disableSpecialEntities {
+                                        incrementLocked(ents, groupMetaKey(k), 1)
+                                }
                                 if old := ents.Get(kb); old != nil {
                                         oldPM, err := rpmWoCtx(old, ns)
                                         if err != nil {
                                                 return err
                                         }
                                         updateIndexes(d.head, k, oldPM, nil)
                                         ents.Delete(kb)
                                         if d.consistent {
                                                 d.snap = d.head.Snapshot()
                                         }
@@ skipping 159 matching lines @@
                 td.muts[rk] = append(td.muts[rk], txnMutation{key, data})
         }
 
         return nil
 }
 
 func (td *txnDataStoreData) putMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.PutMultiCB) {
         ents := td.parent.mutableEnts(keys[0].Namespace())
 
         for i, k := range keys {
-                func() {
+                err := func() (err error) {
                         td.parent.Lock()
                         defer td.parent.Unlock()
-                        k = td.parent.fixKeyLocked(ents, k)
+                        k, err = td.parent.fixKeyLocked(ents, k)
+                        return
                 }()
-                err := td.writeMutation(false, k, vals[i])
+                if err == nil {
+                        err = td.writeMutation(false, k, vals[i])
+                }
                 if cb != nil {
                         cb(k, err)
                 }
         }
 }
 
 func (td *txnDataStoreData) getMulti(keys []*ds.Key, cb ds.GetMultiCB) error {
         return getMultiInner(keys, cb, func() (*memCollection, error) {
                 err := error(nil)
                 for _, key := range keys {
@@ skipping 22 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, "", "")
 }