datastore: Update AllocateIDs to take keys.

impl/memory/datastore_data.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"
         "fmt"
         "strings"
@@ skipping 197 matching lines @@
         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 + int64(amt-1))},
         }))
         return ret
 }
 
-func (d *dataStoreData) allocateIDs(incomplete *ds.Key, n int) (int64, error) {
-        d.Lock()
-        defer d.Unlock()
+func (d *dataStoreData) allocateIDs(keys []*ds.Key, cb ds.NewKeyCB) error {
+        // Map keys by entity type.
+        entityMap := make(map[string][]int)
+        for i, key := range keys {
+                ks := key.String()
+                entityMap[ks] = append(entityMap[ks], i)
+        }
 
-        ents := d.head.GetOrCreateCollection("ents:" + incomplete.Namespace())
-        return d.allocateIDsLocked(ents, incomplete, n)
+        // Allocate IDs for our keys. We use an inline function so we can ensure that
+        // the lock is released.
+        err := func() error {
+                d.Lock()
+                defer d.Unlock()
+
+                for _, idxs := range entityMap {
+                        baseKey := keys[idxs[0]]
+
+                        ents := d.head.GetOrCreateCollection("ents:" + baseKey.Namespace())
+
+                        // Allocate IDs. The only possible error is when disableSpecialEntities is
+                        // true, in which case we will return a full method error instead of
+                        // individual callback errors.
+                        start, err := d.allocateIDsLocked(ents, baseKey, len(idxs))
+                        if err != nil {
+                                return err
+                        }
+
+                        for i, idx := range idxs {
+                                keys[idx] = baseKey.WithID("", start+int64(i))
+                        }
+                }
+                return nil
+        }()
+        if err != nil {
+                return err
+        }
+
+        // Execute Callbacks.
+        for _, key := range keys {
+                cb(key, nil)
+        }
+        return nil
 }
 
 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), nil
 }
 
 func (d *dataStoreData) fixKeyLocked(ents memCollection, key *ds.Key) (*ds.Key, error) {
-        if key.Incomplete() {
+        if key.IsIncomplete() {
                 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, nil
 }
 
-func (d *dataStoreData) putMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.PutMultiCB) error {
+func (d *dataStoreData) putMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.NewKeyCB) error {
         ns := 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.Lock()
                         defer d.Unlock()
 
@@ skipping 248 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) {
+func (td *txnDataStoreData) putMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.NewKeyCB) {
         ns := keys[0].Namespace()
 
         for i, k := range keys {
                 err := func() (err error) {
                         td.parent.Lock()
                         defer td.parent.Unlock()
                         ents := td.parent.head.GetOrCreateCollection("ents:" + ns)
 
                         k, err = td.parent.fixKeyLocked(ents, k)
                         return
@@ skipping 53 matching lines @@
         }
         return namespaces
 }
 
 func trimPrefix(v, p string) (string, bool) {
         if strings.HasPrefix(v, p) {
                 return v[len(p):], true
         }
         return v, false
 }