datastore: Update AllocateIDs to take keys.

filter/txnBuf/state.go

Index: filter/txnBuf/state.go
diff --git a/filter/txnBuf/state.go b/filter/txnBuf/state.go
index 4e22c3d12a584b3aa9b5e1531827852f9210d42b..1e7892f9170fc5b405b26961702598d50200f595 100644
--- a/filter/txnBuf/state.go
+++ b/filter/txnBuf/state.go
@@ -363,30 +363,49 @@ func (t *txnBufState) deleteMulti(keys []*datastore.Key, cb datastore.DeleteMult
 }
 
 func (t *txnBufState) fixKeys(keys []*datastore.Key) ([]*datastore.Key, error) {
-	lme := errors.NewLazyMultiError(len(keys))
-	realKeys := []*datastore.Key(nil)
+	// Identify any incomplete keys and allocate IDs for them.
+	//
+	// In order to facilitate this, we will maintain a mapping of the
+	// incompleteKeys index to the key's corresponding index in the keys array.
+	// Any errors or allocations on incompleteKeys operations will be propagated
+	// to the correct keys index using this map.
+	var incompleteKeys []*datastore.Key
+	incompleteMap := make(map[int]int)
+
 	for i, key := range keys {
 		if key.Incomplete() {
-			// intentionally call AllocateIDs without lock.
-			start, err := t.parentDS.AllocateIDs(key, 1)
-			if !lme.Assign(i, err) {
-				if realKeys == nil {
-					realKeys = make([]*datastore.Key, len(keys))
-					copy(realKeys, keys)
-				}
+			incompleteMap[len(incompleteKeys)] = i
+			incompleteKeys = append(incompleteKeys, key)
+		}
+	}
+	if len(incompleteKeys) == 0 {
+		return keys, nil
+	}
 
-				aid, ns, toks := key.Split()
-				toks[len(toks)-1].IntID = start
-				realKeys[i] = datastore.NewKeyToks(aid, ns, toks)
+	// Intentionally call AllocateIDs without lock.
+	if err := t.parentDS.AllocateIDs(incompleteKeys); err != nil {
+		// If we received a MultiError, map the errors back to our outer
+		// LazyMultiError.
+		if merr, ok := err.(errors.MultiError); ok {
+			outerErr := make(errors.MultiError, len(keys))
+			for i, e := range merr {
+				if e != nil {
+					outerErr[incompleteMap[i]] = e
+				}
 			}
+			err = outerErr
 		}
+		return nil, err
 	}
-	err := lme.Get()
 
-	if realKeys != nil {
-		return realKeys, err
+	// Allocate a new keys slice, copy all initially complete keys into it, and
+	// load the newly-allocated keys into it.
+	keys, origKeys := make([]*datastore.Key, len(keys)), keys
+	copy(keys, origKeys)
+	for i, k := range incompleteKeys {
+		keys[incompleteMap[i]] = k
 	}
-	return keys, err
+	return keys, nil
 }
 
 func (t *txnBufState) putMulti(keys []*datastore.Key, vals []datastore.PropertyMap, cb datastore.PutMultiCB, haveLock bool) error {