datastore: Update AllocateIDs to take keys.

filter/txnBuf/state.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 txnBuf
 
 import (
         "bytes"
         "sync"
 
@@ skipping 345 matching lines @@
         for range keys {
                 if err := cb(nil); err != nil {
                         return err
                 }
         }
 
         return nil
 }
 
 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)
-                                }
-
-                                aid, ns, toks := key.Split()
-                                toks[len(toks)-1].IntID = start
-                                realKeys[i] = datastore.NewKeyToks(aid, ns, toks)
-                        }
+                        incompleteMap[len(incompleteKeys)] = i
+                        incompleteKeys = append(incompleteKeys, key)
                 }
         }
-        err := lme.Get()
+        if len(incompleteKeys) == 0 {
+                return keys, nil
+        }
 
-        if realKeys != nil {
-                return realKeys, err
+        // 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
         }
-        return keys, 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, nil
 }
 
 func (t *txnBufState) putMulti(keys []*datastore.Key, vals []datastore.PropertyMap, cb datastore.PutMultiCB, haveLock bool) error {
         keys, err := t.fixKeys(keys)
         if err != nil {
                 for _, e := range err.(errors.MultiError) {
                         if err := cb(nil, e); err != nil {
                                 return err
                         }
                 }
@@ skipping 142 matching lines @@
 // plus a stringset of all the encoded root keys that `keys` represents.
 func toEncoded(keys []*datastore.Key) (full []string, roots stringset.Set) {
         roots = stringset.New(len(keys))
         full = make([]string, len(keys))
         for i, k := range keys {
                 roots.Add(string(serialize.ToBytes(k.Root())))
                 full[i] = string(serialize.ToBytes(k))
         }
         return
 }