Remove dsQueryBatch in favor of ds.Batcher.

dm/appengine/deps/walk_graph.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 deps
 
 import (
         "fmt"
         "math"
         "sort"
         "sync"
         "time"
 
         "golang.org/x/net/context"
         "google.golang.org/grpc/codes"
 
-        "github.com/luci/gae/filter/dsQueryBatch"
         ds "github.com/luci/gae/service/datastore"
+
         "github.com/luci/luci-go/common/clock"
         "github.com/luci/luci-go/common/logging"
         "github.com/luci/luci-go/common/sync/parallel"
         dm "github.com/luci/luci-go/dm/api/service/v1"
         "github.com/luci/luci-go/dm/appengine/distributor"
         "github.com/luci/luci-go/dm/appengine/model"
         "github.com/luci/luci-go/grpc/grpcutil"
 )
 
 const numWorkers = 16
 
 const maxTimeout = 55 * time.Second // GAE limit is 60s
 
+// queryBatcher is a default datastore Batcher instance for queries.
+//
+// we end up doing a lot of long queries in here, so use a batcher to
+// prevent datastore timeouts.
+var queryBatcher ds.Batcher
+
 type node struct {
         aid                 *dm.Attempt_ID
         depth               int64
         canSeeAttemptResult bool
 }
 
 func (g *graphWalker) runSearchQuery(send func(*dm.Attempt_ID) error, s *dm.GraphQuery_Search) func() error {
         return func() error {
                 logging.Errorf(g, "SearchQuery not implemented")
                 return grpcutil.Errf(codes.Unimplemented, "GraphQuery.Search is not implemented")
         }
 }
 
 func isCtxErr(err error) bool {
         return err == context.Canceled || err == context.DeadlineExceeded
 }
 
 func (g *graphWalker) runAttemptListQuery(send func(*dm.Attempt_ID) error) func() error {
         return func() error {
                 for qst, anum := range g.req.Query.AttemptList.To {
                         if len(anum.Nums) == 0 {
                                 qry := model.QueryAttemptsForQuest(g, qst)
                                 if !g.req.Include.Attempt.Abnormal {
                                         qry = qry.Eq("IsAbnormal", false)
                                 }
                                 if !g.req.Include.Attempt.Expired {
                                         qry = qry.Eq("IsExpired", false)
                                 }
-                                err := ds.Run(g, qry, func(k *ds.Key) error {
+                                err := queryBatcher.Run(g, qry, func(k *ds.Key) error {
                                         aid := &dm.Attempt_ID{}
                                         if err := aid.SetDMEncoded(k.StringID()); err != nil {
                                                 logging.WithError(err).Errorf(g, "Attempt_ID.SetDMEncoded returned an error with input: %q", k.StringID())
                                                 panic(fmt.Errorf("in AttemptListQuery: %s", err))
                                         }
                                         return send(aid)
                                 })
                                 if err != nil {
                                         if !isCtxErr(err) {
                                                 logging.WithError(err).Errorf(g, "in AttemptListQuery")
@@ skipping 35 matching lines @@
                         logging.Fields{ek: err, "qid": qid}.Errorf(g, "while loading quest")
                         return err
                 }
                 dst.Data = qst.DataProto()
                 dst.Partial = false
                 return nil
         }
 }
 
 func (g *graphWalker) loadEdges(send func(*dm.Attempt_ID) error, typ string, base *ds.Key, fan *dm.AttemptList, doSend bool) error {
-        return ds.Run(g, ds.NewQuery(typ).Ancestor(base), func(k *ds.Key) error {
+        return queryBatcher.Run(g, ds.NewQuery(typ).Ancestor(base), func(k *ds.Key) error {
                 if g.Err() != nil {
                         return ds.Stop
                 }
                 aid := &dm.Attempt_ID{}
                 if err := aid.SetDMEncoded(k.StringID()); err != nil {
                         logging.WithError(err).Errorf(g, "Attempt_ID.SetDMEncoded returned an error with input: %q", k.StringID())
                         panic(fmt.Errorf("in AttemptListQuery: %s", err))
                 }
                 if fan != nil {
                         fan.AddAIDs(aid)
@@ skipping 20 matching lines @@
         } else {
                 q = q.Limit(int32(numEx))
         }
         if !g.req.Include.Execution.Abnormal {
                 q = q.Eq("IsAbnormal", false)
         }
         if !g.req.Include.Execution.Expired {
                 q = q.Eq("IsExpired", false)
         }
 
-        return ds.Run(g, q, func(e *model.Execution) error {
+        return queryBatcher.Run(g, q, func(e *model.Execution) error {
                 if g.Err() != nil {
                         return ds.Stop
                 }
                 p := e.ToProto(g.req.Include.Execution.Ids)
                 d, err := g.getDist(p.Data.DistributorInfo.ConfigName)
                 if err != nil {
                         return err
                 }
                 p.Data.DistributorInfo.Url = d.InfoURL(distributor.Token(p.Data.DistributorInfo.Token))
                 if d := p.Data.GetFinished().GetData(); d != nil {
@@ skipping 245 matching lines @@
         return idx < len(nums) && nums[idx] == aid.Id
 }
 
 func doGraphWalk(c context.Context, req *dm.WalkGraphReq) (rsp *dm.GraphData, err error) {
         cncl := (func())(nil)
         timeoutProto := req.Limit.MaxTime
         timeout := timeoutProto.Duration() // .Duration on nil is OK
         if timeoutProto == nil || timeout > maxTimeout {
                 timeout = maxTimeout
         }
-        // we end up doing a lot of long queries in here, so install a batcher to
-        // prevent datastore timeouts.
-        c = dsQueryBatch.BatchQueries(c, 100)
         c, cncl = clock.WithTimeout(c, timeout)
         defer cncl()
 
         // nodeChan recieves attempt nodes to process. If it recieves the
         // `finishedJob` sentinel node, that indicates that an outstanding worker is
         // finished.
         nodeChan := make(chan *node, numWorkers)
         defer close(nodeChan)
 
         g := graphWalker{Context: c, req: req}
@@ skipping 152 matching lines @@
 func (d *deps) WalkGraph(c context.Context, req *dm.WalkGraphReq) (rsp *dm.GraphData, err error) {
         if req.Auth != nil {
                 logging.Fields{"execution": req.Auth.Id}.Debugf(c, "on behalf of")
         } else {
                 if err = canRead(c); err != nil {
                         return
                 }
         }
         return doGraphWalk(c, req)
 }