datastore: Update AllocateIDs to take keys.

impl/memory/datastore_data.go

Index: impl/memory/datastore_data.go
diff --git a/impl/memory/datastore_data.go b/impl/memory/datastore_data.go
index 94a868878e6db035cb3de4e9b4e6758e8e514e2f..4551c4bef3611e25da8e5c0ddedf82c28cf8d17d 100644
--- a/impl/memory/datastore_data.go
+++ b/impl/memory/datastore_data.go
@@ -215,17 +215,53 @@ func incrementLocked(ents memCollection, key []byte, amt int) int64 {
 	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)
+	}
+
+	// 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
+			}
 
-	ents := d.head.GetOrCreateCollection("ents:" + incomplete.Namespace())
-	return d.allocateIDsLocked(ents, incomplete, n)
+			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)