Interface update, per-method Contexts.

impl/memory/context.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 (
         "errors"
         "strings"
         "sync"
 
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/luci-go/common/logging/memlogger"
+
         "golang.org/x/net/context"
 )
 
 var serializationDeterministic = false
 
 type memContextObj interface {
         sync.Locker
         canApplyTxn(m memContextObj) bool
         applyTxn(c context.Context, m memContextObj)
 
@@ skipping 97 matching lines @@
         }
         c = memlogger.Use(c)
 
         fqAppID := aid
         if parts := strings.SplitN(fqAppID, "~", 2); len(parts) == 2 {
                 aid = parts[1]
         }
 
         memctx := newMemContext(fqAppID)
         c = context.WithValue(c, memContextKey, memctx)
-        c = context.WithValue(c, memContextNoTxnKey, memctx)
         c = useGID(c, func(mod *globalInfoData) {
                 mod.appID = aid
                 mod.fqAppID = fqAppID
         })
         return useMod(useMail(useUser(useTQ(useRDS(useMC(useGI(c)))))))
 }
 
-func cur(c context.Context) (p *memContext) {
-        p, _ = c.Value(memContextKey).(*memContext)
-        return
-}
-
-func curNoTxn(c context.Context) (p *memContext) {
-        p, _ = c.Value(memContextNoTxnKey).(*memContext)
-        return
+func cur(c context.Context) (*memContext, bool) {
+        if txn := c.Value(currentTxnKey); txn != nil {
+                // We are in a Transaction.
+                return txn.(*memContext), true
+        }
+        return c.Value(memContextKey).(*memContext), false
 }
 
 type memContextKeyType int
 
 var (
-        memContextKey      memContextKeyType
-        memContextNoTxnKey memContextKeyType = 1
+        memContextKey memContextKeyType
+        currentTxnKey = 1
 )
 
 // weird stuff
 
 // RunInTransaction is here because it's really a service-wide transaction, not
 // just in the datastore. TaskQueue behaves differently in a transaction in
 // a couple ways, for example.
 //
 // It really should have been appengine.Context.RunInTransaction(func(tc...)),
 // but because it's not, this method is on dsImpl instead to mirror the official
 // API.
 //
 // The fake implementation also differs from the real implementation because the
 // fake TaskQueue is NOT backed by the fake Datastore. This is done to make the
 // test-access API for TaskQueue better (instead of trying to reconstitute the
 // state of the task queue from a bunch of datastore accesses).
 func (d *dsImpl) RunInTransaction(f func(context.Context) error, o *ds.TransactionOptions) error {
         if d.data.getDisableSpecialEntities() {
                 return errors.New("special entities are disabled. no transactions for you")
         }
 
         // Keep in separate function for defers.
         loopBody := func(applyForReal bool) error {
-                curMC := cur(d.c)
+                curMC, inTxn := cur(d)
+                if inTxn {
+                        return errors.New("datastore: nested transactions are not supported")
+                }
 
                 txnMC := curMC.mkTxn(o)
 
                 defer func() {
                         txnMC.Lock()
                         defer txnMC.Unlock()
 
                         txnMC.endTxn()
                 }()
 
-                if err := f(context.WithValue(d.c, memContextKey, txnMC)); err != nil {
+                if err := f(context.WithValue(d, currentTxnKey, txnMC)); err != nil {
                         return err
                 }
 
                 txnMC.Lock()
                 defer txnMC.Unlock()
 
                 if applyForReal && curMC.canApplyTxn(txnMC) {
-                        curMC.applyTxn(d.c, txnMC)
+                        curMC.applyTxn(d, txnMC)
                 } else {
                         return ds.ErrConcurrentTransaction
                 }
                 return nil
         }
 
         // From GAE docs for TransactionOptions: "If omitted, it defaults to 3."
         attempts := 3
         if o != nil && o.Attempts != 0 {
                 attempts = o.Attempts
         }
         for attempt := 0; attempt < attempts; attempt++ {
                 if err := loopBody(attempt >= d.data.txnFakeRetry); err != ds.ErrConcurrentTransaction {
                         return err
                 }
         }
         return ds.ErrConcurrentTransaction
 }