datastore: Update AllocateIDs to take keys.

filter/txnBuf/state.go

 // Copyright 2015 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // 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  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)
+                if key.IsIncomplete() {
+                        if incompleteMap == nil {
+                                incompleteMap = make(map[int]int)
                         }
+                        incompleteMap[len(incompleteKeys)] = i
+                        incompleteKeys = append(incompleteKeys, key)
                 }
         }
-        err := lme.Get()
+        if len(incompleteKeys) == 0 {
+                return keys, nil
+        }
 
-        if realKeys != nil {
-                return realKeys, err
+        // We're going to update keys, so clone it.
+        keys, origKeys := make([]*datastore.Key, len(keys)), keys
+        copy(keys, origKeys)
+
+        // Intentionally call AllocateIDs without lock.
+        i := 0
+        outerErr := errors.NewLazyMultiError(len(keys))
+        err := t.parentDS.AllocateIDs(incompleteKeys, func(key *datastore.Key, err error) error {
+                outerIdx := incompleteMap[i]
+                i++
+
+                if err != nil {
+                        outerErr.Assign(outerIdx, err)
+                } else {
+                        keys[outerIdx] = key
+                }
+                return nil
+        })
+        if err != nil {
+                return nil, err
         }
-        return keys, err
+        return keys, outerErr.Get()
 }
 
-func (t *txnBufState) putMulti(keys []*datastore.Key, vals []datastore.PropertyMap, cb datastore.PutMultiCB, haveLock bool) error {
+func (t *txnBufState) putMulti(keys []*datastore.Key, vals []datastore.PropertyMap, cb datastore.NewKeyCB, 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
                         }
                 }
                 return nil
         }
 
@@ skipping 139 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
 }