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) error {
+        start, err := func() (int64, error) {
+                d.Lock()
+                defer d.Unlock()
 
-        ents := d.head.GetOrCreateCollection("ents:" + incomplete.Namespace())
-        return d.allocateIDsLocked(ents, incomplete, n)
+                // Allocate a set of IDs in our key namespace. We know that "keys" has at
+                // least one entry and that all keys share the same namespace and entty
+                // type because checkfilter has already asserted this.
+                baseKey := keys[0]
+                ents := d.head.GetOrCreateCollection("ents:" + baseKey.Namespace())
+                start, err := d.allocateIDsLocked(ents, baseKey, len(keys))
+                if err != nil {
+                        return -1, err
+                }
+                return start, nil
+        }()
+        if err != nil {
+                return err
+        }
+
+        // Update the keys in our "keys" slice with their new IDs.
+        for i, k := range keys {
+                keys[i] = ds.NewKey(k.AppID(), k.Namespace(), k.Kind(), "", start+int64(i), k.Parent())
+        }
+        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())
@@ skipping 350 matching lines @@
         }
         return namespaces
 }
 
 func trimPrefix(v, p string) (string, bool) {
         if strings.HasPrefix(v, p) {
                 return v[len(p):], true
         }
         return v, false
 }