datastore: Update AllocateIDs to take keys.

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"
         "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.PutMultiCB) 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")
+                return -1, errors.New("disableSpecialEntities is true so allocateIDs is disabled")

[iannucci, 2016/06/14 01:46:04]

0 is an invalid id, this shouldn't need to be -1.

[dnj (Google), 2016/06/14 04:24:18]

Done.

         }
 
         idKey := []byte(nil)
         if incomplete.Parent() == nil {
                 idKey = rootIDsKey(incomplete.Kind())
         } else {
                 idKey = groupIDsKey(incomplete)
         }
         return incrementLocked(ents, idKey, n), nil
 }
@@ skipping 345 matching lines @@
         }
         return namespaces
 }
 
 func trimPrefix(v, p string) (string, bool) {
         if strings.HasPrefix(v, p) {
                 return v[len(p):], true
         }
         return v, false
 }