Add context-aware "time" library wrapper.

go/src/infra/gae/libs/wrapper/memory/taskqueue.go

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package memory
 
 import (
         "fmt"
         "infra/gae/libs/wrapper"
-        "math/rand"
         "net/http"
         "regexp"
-        "time"
 
         "golang.org/x/net/context"
 
         "appengine"
         "appengine/taskqueue"
         "appengine_internal"
         dbpb "appengine_internal/datastore"
         pb "appengine_internal/taskqueue"
 )
 
 /////////////////////////////// public functions ///////////////////////////////
 
 func useTQ(c context.Context) context.Context {
         return wrapper.SetTQFactory(c, func(ic context.Context) wrapper.TaskQueue {
                 tqd := cur(ic).Get(memContextTQIdx)
                 var ret interface {
                         wrapper.TQTestable
                         wrapper.TaskQueue
                 }
                 switch x := tqd.(type) {
                 case *taskQueueData:
                         ret = &taskqueueImpl{
                                 wrapper.DummyTQ(),
                                 x,
+                                ic,
                                 curGID(ic).namespace,
-                                func() time.Time { return wrapper.GetTimeNow(ic) },
-                                wrapper.GetMathRand(ic),
                         }
 
                 case *txnTaskQueueData:
                         ret = &taskqueueTxnImpl{
                                 wrapper.DummyTQ(),
                                 x,
+                                ic,
                                 curGID(ic).namespace,
-                                func() time.Time { return wrapper.GetTimeNow(ic) },
-                                wrapper.GetMathRand(ic),
                         }
 
                 default:
                         panic(fmt.Errorf("TQ: bad type: %v", tqd))
                 }
                 return ret
         })
 }
 
 //////////////////////////////// taskqueueImpl /////////////////////////////////
 
 type taskqueueImpl struct {
         wrapper.TaskQueue
         *taskQueueData
 
-        ns       string
-        timeNow  func() time.Time
-        mathRand *rand.Rand
+        ctx context.Context
+        ns  string
 }
 
 var (
         _ = wrapper.TaskQueue((*taskqueueImpl)(nil))
         _ = wrapper.TQTestable((*taskqueueImpl)(nil))
 )
 
 func (t *taskqueueImpl) addLocked(task *taskqueue.Task, queueName string) (*taskqueue.Task, error) {
-        toSched, queueName, err := t.prepTask(t.ns, task, queueName, t.timeNow(), t.mathRand)
+        toSched, queueName, err := t.prepTask(t.ctx, t.ns, task, queueName)
         if err != nil {
                 return nil, err
         }
 
         if _, ok := t.archived[queueName][toSched.Name]; ok {
                 // SDK converts TOMBSTONE -> already added too
                 return nil, taskqueue.ErrTaskAlreadyAdded
         } else if _, ok := t.named[queueName][toSched.Name]; ok {
                 return nil, taskqueue.ErrTaskAlreadyAdded
         } else {
@@ skipping 70 matching lines @@
                 })
         return err
 }
 
 /////////////////////////////// taskqueueTxnImpl ///////////////////////////////
 
 type taskqueueTxnImpl struct {
         wrapper.TaskQueue
         *txnTaskQueueData
 
-        ns       string
-        timeNow  func() time.Time
-        mathRand *rand.Rand
+        ctx context.Context
+        ns  string
 }
 
 var (
         _ = wrapper.TaskQueue((*taskqueueTxnImpl)(nil))
         _ = wrapper.TQTestable((*taskqueueTxnImpl)(nil))
 )
 
 func (t *taskqueueTxnImpl) addLocked(task *taskqueue.Task, queueName string) (*taskqueue.Task, error) {
-        toSched, queueName, err := t.parent.prepTask(t.ns, task, queueName, t.timeNow(), t.mathRand)
+        toSched, queueName, err := t.parent.prepTask(t.ctx, t.ns, task, queueName)
         if err != nil {
                 return nil, err
         }
 
         numTasks := 0
         for _, vs := range t.anony {
                 numTasks += len(vs)
         }
         if numTasks+1 > 5 {
                 // transactional tasks are actually implemented 'for real' as Actions which
@@ skipping 38 matching lines @@
 
         return multi(tasks, queueName, t.addLocked)
 }
 
 ////////////////////////////// private functions ///////////////////////////////
 
 var validTaskName = regexp.MustCompile("^[0-9a-zA-Z\\-\\_]{0,500}$")
 
 const validTaskChars = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ-_"
 
-func mkName(rnd *rand.Rand, cur string, queue map[string]*taskqueue.Task) string {
+func mkName(c context.Context, cur string, queue map[string]*taskqueue.Task) string {
         _, ok := queue[cur]
         for !ok && cur == "" {
                 name := [500]byte{}
                 for i := 0; i < 500; i++ {
-                        name[i] = validTaskChars[rnd.Intn(len(validTaskChars))]
+                        name[i] = validTaskChars[wrapper.GetMathRand(c).Intn(len(validTaskChars))]
                 }
                 cur = string(name[:])
                 _, ok = queue[cur]
         }
         return cur
 }
 
 func newTQError(code pb.TaskQueueServiceError_ErrorCode) *appengine_internal.APIError {
         return &appengine_internal.APIError{Service: "taskqueue", Code: int32(code)}
 }
@@ skipping 45 matching lines @@
 func dupQueue(q wrapper.QueueData) wrapper.QueueData {
         r := make(wrapper.QueueData, len(q))
         for k, q := range q {
                 r[k] = make(map[string]*taskqueue.Task, len(q))
                 for tn, t := range q {
                         r[k][tn] = dupTask(t)
                 }
         }
         return r
 }