datastore: Update AllocateIDs to take keys.

filter/featureBreaker/rds.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 featureBreaker
 
 import (
         "golang.org/x/net/context"
 
         ds "github.com/luci/gae/service/datastore"
 )
 
 type dsState struct {
         *state
 
         rds ds.RawInterface
 }
 
-func (r *dsState) AllocateIDs(incomplete *ds.Key, n int) (int64, error) {
-        start := int64(0)
-        err := r.run(func() (err error) {
-                start, err = r.rds.AllocateIDs(incomplete, n)
-                return
+func (r *dsState) AllocateIDs(keys []*ds.Key, cb ds.NewKeyCB) error {
+        return r.run(func() error {
+                return r.rds.AllocateIDs(keys, cb)
         })
-        return start, err
 }
 
 func (r *dsState) DecodeCursor(s string) (ds.Cursor, error) {
         curs := ds.Cursor(nil)
         err := r.run(func() (err error) {
                 curs, err = r.rds.DecodeCursor(s)
                 return
         })
         return curs, err
 }
@@ skipping 27 matching lines @@
                 return r.rds.DeleteMulti(keys, cb)
         })
 }
 
 func (r *dsState) GetMulti(keys []*ds.Key, meta ds.MultiMetaGetter, cb ds.GetMultiCB) error {
         return r.run(func() error {
                 return r.rds.GetMulti(keys, meta, cb)
         })
 }
 
-func (r *dsState) PutMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.PutMultiCB) error {
+func (r *dsState) PutMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.NewKeyCB) error {
         return r.run(func() (err error) {
                 return r.rds.PutMulti(keys, vals, cb)
         })
 }
 
 func (r *dsState) Testable() ds.Testable {
         return r.rds.Testable()
 }
 
 // FilterRDS installs a featureBreaker datastore filter in the context.
 func FilterRDS(c context.Context, defaultError error) (context.Context, FeatureBreaker) {
         state := newState(defaultError)
         return ds.AddRawFilters(c, func(ic context.Context, RawDatastore ds.RawInterface) ds.RawInterface {
                 return &dsState{state, RawDatastore}
         }), state
 }