Add missing AllocateIDs API.

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 91 matching lines @@
                         if ok && len(pl) > 0 && pl[0].Type() == ds.PTInt {
                                 return pl[0].Value().(int64)
                         }
 
                         memoryCorruption(fmt.Errorf("__version__ property missing or wrong: %v", pm))
                 }
         }
         return 0
 }
 
-func incrementLocked(ents *memCollection, key []byte) int64 {
+func incrementLocked(ents *memCollection, key []byte, amt int) int64 {
+        if amt <= 0 {
+                panic(fmt.Errorf("incrementLocked called with bad `amt`: %d", amt))
+        }
         ret := curVersion(ents, key) + 1
         ents.Set(key, serialize.ToBytes(ds.PropertyMap{
-                "__version__": {ds.MkPropertyNI(ret)},
+                "__version__": {ds.MkPropertyNI(ret + int64(amt-1))},
         }))
         return ret
 }
 
-func (d *dataStoreData) entsKeyLocked(key *ds.Key) (*memCollection, *ds.Key) {
-        coll := "ents:" + key.Namespace()
+func (d *dataStoreData) mutableEnts(ns string) *memCollection {
+        d.Lock()
+        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 {
+        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 {
+        idKey := []byte(nil)
+        if incomplete.Parent() == nil {
+                idKey = rootIDsKey(incomplete.Last().Kind)
+        } else {
+                idKey = groupIDsKey(incomplete)
+        }
+        return incrementLocked(ents, idKey, n)
+}
+
+func (d *dataStoreData) fixKeyLocked(ents *memCollection, key *ds.Key) *ds.Key {
         if key.Incomplete() {
-                idKey := []byte(nil)
-                if key.Parent() == nil {
-                        idKey = rootIDsKey(key.Last().Kind)
-                } else {
-                        idKey = groupIDsKey(key)
-                }
-                id := incrementLocked(ents, idKey)
+                id := d.allocateIDsLocked(ents, key, 1)
                 key = ds.NewKey(key.AppID(), key.Namespace(), key.Last().Kind, "", id, key.Parent())
         }
-
-        return ents, key
+        return key
 }
 
 func (d *dataStoreData) putMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.PutMultiCB) {
+        ents := d.mutableEnts(keys[0].Namespace())
+
         for i, k := range keys {
                 pmap, _ := vals[i].Save(false)
                 dataBytes := serialize.ToBytes(pmap)
 
                 k, err := func() (ret *ds.Key, err error) {
-                        d.rwlock.Lock()
-                        defer d.rwlock.Unlock()
+                        d.Lock()
+                        defer d.Unlock()
 
-                        ents, ret := d.entsKeyLocked(k)
-                        incrementLocked(ents, groupMetaKey(ret))
+                        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 {
                                         return
                                 }
                         }
                         updateIndexes(d.head, ret, oldPM, pmap)
                         ents.Set(keyBytes(ret), dataBytes)
@@ skipping 43 matching lines @@
         }
         ns := keys[0].Namespace()
 
         d.rwlock.Lock()
         defer d.rwlock.Unlock()
 
         ents := d.head.GetCollection("ents:" + ns)
 
         for i, k := range keys {
                 if ents != nil {
-                        incrementLocked(ents, groupMetaKey(k))
+                        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)
                                         }
                                         continue
                                 }
                                 updateIndexes(d.head, k, oldPM, nil)
@@ skipping 151 matching lines @@
                 td.muts[rk] = []txnMutation{}
         }
         if !getOnly {
                 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() {
                         td.parent.Lock()
                         defer td.parent.Unlock()
-                        _, k = td.parent.entsKeyLocked(k)
+                        k = td.parent.fixKeyLocked(ents, k)
                 }()
                 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) {
@@ skipping 24 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, "", "")
 }