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"
+        "infra/libs/clock"
         "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
 
+        ctx      context.Context
         ns       string
-        timeNow  func() time.Time
         mathRand *rand.Rand

[iannucci, 2015/06/03 17:37:19]

could we pull out mathRand the same way too? (e.g. delete and just call
wrapper.MathRand on ctx)

[dnj, 2015/06/03 18:21:26]

Done.

 }
 
 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.ns, task, queueName, clock.Now(t.ctx), t.mathRand)
         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
 
+        ctx      context.Context
         ns       string
-        timeNow  func() time.Time
         mathRand *rand.Rand
 }
 
 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.ns, task, queueName, clock.Now(t.ctx), t.mathRand)
         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 109 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
 }