Interface update, per-method Contexts.

impl/memory/taskqueue.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 memory
 
 import (
         "regexp"
-        "sync/atomic"
 
         "golang.org/x/net/context"
 
         tq "github.com/luci/gae/service/taskqueue"
+
         "github.com/luci/luci-go/common/data/rand/mathrand"
         "github.com/luci/luci-go/common/errors"
 )
 
 /////////////////////////////// public functions ///////////////////////////////
 
 func useTQ(c context.Context) context.Context {
-        return tq.SetRawFactory(c, func(ic context.Context, wantTxn bool) tq.RawInterface {
-                ns, _ := curGID(ic).getNamespace()
-                var tqd memContextObj
+        return tq.SetRawFactory(c, func(ic context.Context) tq.RawInterface {
+                memCtx, isTxn := cur(ic)
+                tqd := memCtx.Get(memContextTQIdx)
 
-                if !wantTxn {
-                        tqd = curNoTxn(ic).Get(memContextTQIdx)
-                } else {
-                        tqd = cur(ic).Get(memContextTQIdx)
+                if isTxn {
+                        return &taskqueueTxnImpl{tqd.(*txnTaskQueueData), ic}
                 }
-
-                if x, ok := tqd.(*taskQueueData); ok {
-                        return &taskqueueImpl{x, ic, ns}
-                }
-                return &taskqueueTxnImpl{tqd.(*txnTaskQueueData), ic, ns}
+                return &taskqueueImpl{tqd.(*taskQueueData), ic}
         })
 }
 
 //////////////////////////////// taskqueueImpl /////////////////////////////////
 
 type taskqueueImpl struct {
         *taskQueueData
 
         ctx context.Context
-        ns  string
 }
 
 var (
         _ = tq.RawInterface((*taskqueueImpl)(nil))
         _ = tq.Testable((*taskqueueImpl)(nil))
 )
 
 func (t *taskqueueImpl) addLocked(task *tq.Task, queueName string) (*tq.Task, error) {
-        toSched, err := t.prepTask(t.ctx, t.ns, task, queueName)
+        toSched, err := t.prepTask(t.ctx, task, queueName)
         if err != nil {
                 return nil, err
         }
 
         if _, ok := t.archived[queueName][toSched.Name]; ok {
                 // SDK converts TOMBSTONE -> already added too
                 return nil, tq.ErrTaskAlreadyAdded
         } else if _, ok := t.named[queueName][toSched.Name]; ok {
                 return nil, tq.ErrTaskAlreadyAdded
         } else {
@@ skipping 74 matching lines @@
                                         s.OldestETA = t.ETA
                                 }
                         }
                         cb(&s, nil)
                 }
         }
 
         return nil
 }
 
-func (t *taskqueueImpl) Testable() tq.Testable {
-        return t
-}
+func (t *taskqueueImpl) GetTestable() tq.Testable { return t }
 
 /////////////////////////////// taskqueueTxnImpl ///////////////////////////////
 
 type taskqueueTxnImpl struct {
         *txnTaskQueueData
 
         ctx context.Context
-        ns  string
 }
 
 var _ interface {
         tq.RawInterface
         tq.Testable
 } = (*taskqueueTxnImpl)(nil)
 
 func (t *taskqueueTxnImpl) addLocked(task *tq.Task, queueName string) (*tq.Task, error) {
-        toSched, err := t.parent.prepTask(t.ctx, t.ns, task, queueName)
+        toSched, err := t.parent.prepTask(t.ctx, 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 11 matching lines @@
         //              We should verify that the .Name for a task added in a transaction is
         //              meaningless. Maybe names generated in a transaction are somehow
         //              guaranteed to be meaningful?
         toRet := toSched.Duplicate()
         toRet.Name = ""
 
         return toRet, nil
 }
 
 func (t *taskqueueTxnImpl) AddMulti(tasks []*tq.Task, queueName string, cb tq.RawTaskCB) error {
-        if atomic.LoadInt32(&t.closed) == 1 {
-                return errors.New("taskqueue: transaction context has expired")
+        if err := assertTxnValid(t.ctx); err != nil {
+                return err
         }
 
         t.Lock()
         defer t.Unlock()
 
         queueName, err := t.parent.getQueueNameLocked(queueName)
         if err != nil {
                 return err
         }
 
         for _, task := range tasks {
                 cb(t.addLocked(task, queueName))
         }
         return nil
 }
 
 func (t *taskqueueTxnImpl) DeleteMulti([]*tq.Task, string, tq.RawCB) error {
         return errors.New("taskqueue: cannot DeleteMulti from a transaction")
 }
 
 func (t *taskqueueTxnImpl) Purge(string) error {
         return errors.New("taskqueue: cannot Purge from a transaction")
 }
 
 func (t *taskqueueTxnImpl) Stats([]string, tq.RawStatsCB) error {
         return errors.New("taskqueue: cannot Stats from a transaction")
 }
 
-func (t *taskqueueTxnImpl) Testable() tq.Testable {
-        return t
-}
+func (t *taskqueueTxnImpl) GetTestable() tq.Testable { return t }
 
 ////////////////////////////// private functions ///////////////////////////////
 
 var validTaskName = regexp.MustCompile("^[0-9a-zA-Z\\-\\_]{0,500}$")
 
 const validTaskChars = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ-_"
 
 func mkName(c context.Context, cur string, queue map[string]*tq.Task) string {
         _, ok := queue[cur]
         for !ok && cur == "" {
@@ skipping 10 matching lines @@
 func dupQueue(q tq.QueueData) tq.QueueData {
         r := make(tq.QueueData, len(q))
         for k, q := range q {
                 r[k] = make(map[string]*tq.Task, len(q))
                 for tn, t := range q {
                         r[k][tn] = t.Duplicate()
                 }
         }
         return r
 }