datastore: Update AllocateIDs to take keys.

impl/prod/raw_datastore.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 prod
 
 import (
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/luci-go/common/errors"
         "golang.org/x/net/context"
@@ skipping 39 matching lines @@
         me, ok := err.(errors.MultiError)
         if ok {
                 for i, err := range me {
                         cb(i, err)
                 }
                 return nil
         }
         return err
 }
 
-func (d rdsImpl) AllocateIDs(incomplete *ds.Key, n int) (start int64, err error) {
-        par, err := dsF2R(d.aeCtx, incomplete.Parent())
-        if err != nil {
-                return
+func (d rdsImpl) AllocateIDs(keys []*ds.Key, cb ds.NewKeyCB) 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)
         }
 
-        start, _, err = datastore.AllocateIDs(d.aeCtx, incomplete.Kind(), par, n)
-        return
+        // Allocate a set of IDs for each unique entity type.
+        errors := errors.NewLazyMultiError(len(keys))
+        setErrs := func(idxs []int, err error) {
+                for _, idx := range idxs {
+                        errors.Assign(idx, err)
+                }
+        }
+
+        for _, idxs := range entityMap {
+                incomplete := keys[idxs[0]]
+                par, err := dsF2R(d.aeCtx, incomplete.Parent())
+                if err != nil {
+                        setErrs(idxs, err)
+                        continue
+                }
+
+                start, _, err := datastore.AllocateIDs(d.aeCtx, incomplete.Kind(), par, len(idxs))
+                if err != nil {
+                        setErrs(idxs, err)
+                        continue
+                }
+
+                for i, idx := range idxs {
+                        keys[idx] = incomplete.WithID("", start+int64(i))
+                }
+        }
+
+        for i, key := range keys {
+                if err := errors.GetOne(i); err != nil {
+                        cb(nil, err)
+                } else {
+                        cb(key, nil)
+                }
+        }
+        return nil
 }
 
 func (d rdsImpl) DeleteMulti(ks []*ds.Key, cb ds.DeleteMultiCB) error {
         keys, err := dsMF2R(d.aeCtx, ks)
         if err == nil {
                 err = datastore.DeleteMulti(d.aeCtx, keys)
         }
         return idxCallbacker(err, len(ks), func(_ int, err error) {
                 cb(err)
         })
@@ skipping 10 matching lines @@
         }
         return idxCallbacker(err, len(keys), func(idx int, err error) {
                 if pls := vals[idx]; pls != nil {
                         cb(pls.(*typeFilter).pm, err)
                 } else {
                         cb(nil, err)
                 }
         })
 }
 
-func (d rdsImpl) PutMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.PutMultiCB) error {
+func (d rdsImpl) PutMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.NewKeyCB) error {
         rkeys, err := dsMF2R(d.aeCtx, keys)
         if err == nil {
                 rvals := make([]datastore.PropertyLoadSaver, len(vals))
                 for i, val := range vals {
                         rvals[i] = &typeFilter{d.aeCtx, val}
                 }
                 rkeys, err = datastore.PutMulti(d.aeCtx, rkeys, rvals)
         }
         return idxCallbacker(err, len(keys), func(idx int, err error) {
                 k := (*ds.Key)(nil)
@@ skipping 121 matching lines @@
 func (d rdsImpl) RunInTransaction(f func(c context.Context) error, opts *ds.TransactionOptions) error {
         ropts := (*datastore.TransactionOptions)(opts)
         return datastore.RunInTransaction(d.aeCtx, func(c context.Context) error {
                 return f(context.WithValue(d.userCtx, prodContextKey, c))
         }, ropts)
 }
 
 func (d rdsImpl) Testable() ds.Testable {
         return nil
 }