Remove dsQueryBatch in favor of ds.Batcher.

appengine/tsmon/handler.go

 // Copyright 2016 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 tsmon
 
 import (
         "net/http"
 
         "golang.org/x/net/context"
 
-        "github.com/luci/gae/filter/dsQueryBatch"
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/gae/service/info"
 
         "github.com/luci/luci-go/common/clock"
         "github.com/luci/luci-go/common/logging"
         "github.com/luci/luci-go/common/tsmon"
         "github.com/luci/luci-go/server/router"
 )
 
-const housekeepingInstanceBatchSize = 200
-
 // InstallHandlers installs HTTP handlers for tsmon routes.
 func InstallHandlers(r *router.Router, base router.MiddlewareChain) {
         r.GET("/internal/cron/ts_mon/housekeeping", base, housekeepingHandler)
 }
 
 // housekeepingHandler is an HTTP handler that should be run every minute by
 // cron on App Engine.  It assigns task numbers to datastore entries, and runs
 // any global metric callbacks.
 func housekeepingHandler(c *router.Context) {
         if !info.IsDevAppServer(c.Context) && c.Request.Header.Get("X-Appengine-Cron") != "true" {
@@ skipping 14 matching lines @@
 // and expiring old entities.
 func assignTaskNumbers(c context.Context) error {
         c = info.MustNamespace(c, instanceNamespace)
 
         now := clock.Now(c)
         expiredTime := now.Add(-instanceExpirationTimeout)
 
         usedTaskNums := map[int]struct{}{}
         totalExpired := 0
 
-        expiredKeys := make([]*ds.Key, 0, housekeepingInstanceBatchSize)
+        expiredKeys := make([]*ds.Key, 0, ds.Raw(c).Constraints().QueryBatchSize)
         var unassigned []*instance
 
         // expireInstanceBatch processes the set of instances in "expiredKeys",
         // deletes them, and clears the list for the next iteration.
         //
         // We do this in batches to handle large numbers without inflating memory
         // requirements. If there are any timeouts or problems, this will also enable
         // us to iteratively chip away at the problem.
         expireInstanceBatch := func(c context.Context) error {
                 if len(expiredKeys) == 0 {
                         return nil
                 }
 
                 logging.Debugf(c, "Expiring %d instance(s)", len(expiredKeys))
                 if err := ds.Delete(c, expiredKeys); err != nil {
                         logging.WithError(err).Errorf(c, "Failed to expire instances.")
                         return err
                 }
 
                 // Clear the instances list for next round.
                 totalExpired += len(expiredKeys)
                 expiredKeys = expiredKeys[:0]
                 return nil
         }
 
         // Query all instances from datastore.
         q := ds.NewQuery("Instance")
-        if err := ds.Run(dsQueryBatch.BatchQueries(c, housekeepingInstanceBatchSize, expireInstanceBatch), q, func(i *instance) {
+
+        b := ds.Batcher{
+                Callback: expireInstanceBatch,
+        }
+        if err := b.Run(c, q, func(i *instance) {
                 if i.TaskNum >= 0 {
                         usedTaskNums[i.TaskNum] = struct{}{}
                 }
                 if i.LastUpdated.Before(expiredTime) {
                         expiredKeys = append(expiredKeys, ds.NewKey(c, "Instance", i.ID, 0, nil))
                         logging.Debugf(c, "Expiring %s task_num %d, inactive since %s",
                                 i.ID, i.TaskNum, i.LastUpdated.String())
                 } else if i.TaskNum < 0 {
                         unassigned = append(unassigned, i)
                 }
@@ skipping 32 matching lines @@
                 for {
                         n := next
                         next++
                         _, has := used[n]
                         if !has {
                                 return n
                         }
                 }
         }
 }