Interface update, per-method Contexts.

impl/prod/context.go

Index: impl/prod/context.go
diff --git a/impl/prod/context.go b/impl/prod/context.go
index fc4a5e8fd9afc7036620e605b4f60fb2f3b85878..de55d5c11a91981710189a929d5ab6402d0bb59b 100644
--- a/impl/prod/context.go
+++ b/impl/prod/context.go
@@ -11,7 +11,6 @@ import (
 	"net/url"
 	"strings"
 
-	"github.com/luci/gae/service/info"
 	"github.com/luci/gae/service/urlfetch"
 	"golang.org/x/net/context"
 	gOAuth "golang.org/x/oauth2/google"
@@ -29,9 +28,8 @@ var RemoteAPIScopes = []string{
 type key int
 
 var (
-	prodContextKey      key

[dnj, 2016/09/01 15:25:40]

This is a dangerous/big change, only b/c it's hard to test this stuff live.
Basically, now instead of tracking the w/ and w/out Transaction Contexts
separately, we store one "prodState" object (see below). That object maintains a
"current" Context and the latest "No Transaction" Context, and sets "Current" to
"No Transaction" when WithTransaction is invoked with nil.

Info service namespace operations install a new prodState with both Contexts
updated to the new namespace.

Consequently, AEContextNoTxn need not exist; one just passes WithTransaction(c,
nil) to AEContext.

-	prodContextNoTxnKey key = 1
-	probeCacheKey       key = 2
+	prodStateKey  key

[iannucci, 2016/09/16 01:01:13]

lets change these to 

string prodStateKey = "contains a *Type"


WithValue(&prodStateKey, ...)

[dnj, 2016/09/16 05:44:43]

Done.

+	probeCacheKey key = 1
 )
 
 // AEContext retrieves the raw "google.golang.org/appengine" compatible Context.
@@ -40,40 +38,15 @@ var (
 // RPCs. Doesn't transfer cancelation ability though (since it's ignored by GAE
 // anyway).
 func AEContext(c context.Context) context.Context {
-	aeCtx, _ := c.Value(prodContextKey).(context.Context)
-	if aeCtx == nil {
-		return nil
-	}
-	if deadline, ok := c.Deadline(); ok {
-		aeCtx, _ = context.WithDeadline(aeCtx, deadline)
-	}
-	return aeCtx
-}
-
-// AEContextNoTxn retrieves the raw "google.golang.org/appengine" compatible
-// Context that's not part of a transaction.
-func AEContextNoTxn(c context.Context) context.Context {
-	aeCtx, _ := c.Value(prodContextNoTxnKey).(context.Context)
-	if aeCtx == nil {
-		return nil
-	}
-
-	if ns, has := info.Get(c).GetNamespace(); has {
-		var err error
-		aeCtx, err = appengine.Namespace(aeCtx, ns)
-		if err != nil {
-			panic(err)
-		}
-	}
-	if deadline, ok := c.Deadline(); ok {
-		aeCtx, _ = context.WithDeadline(aeCtx, deadline)
-	}
-	return aeCtx
+	ps := getProdState(c)
+	return ps.context(c)
 }
 
 func setupAECtx(c, aeCtx context.Context) context.Context {
-	c = context.WithValue(c, prodContextKey, aeCtx)
-	c = context.WithValue(c, prodContextNoTxnKey, aeCtx)
+	c = withProdState(c, prodState{
+		ctx:      aeCtx,
+		noTxnCtx: aeCtx,
+	})
 	return useModule(useMail(useUser(useURLFetch(useRDS(useMC(useTQ(useGI(useLogging(c)))))))))
 }
 
@@ -123,9 +96,11 @@ func Use(c context.Context, r *http.Request) context.Context {
 //       - "https://www.googleapis.com/auth/cloud.platform"
 func UseRemote(inOutCtx *context.Context, host string, client *http.Client) (err error) {
 	if client == nil {
+		aeCtx := AEContext(*inOutCtx)

[dnj, 2016/09/01 15:25:40]

As far as I can tell, the NoTxn part of the previous code was not important
here.

+
 		if strings.HasPrefix(host, "localhost") {
 			transp := http.DefaultTransport
-			if aeCtx := AEContextNoTxn(*inOutCtx); aeCtx != nil {
+			if aeCtx != nil {
 				transp = urlfetch.Get(*inOutCtx)
 			}
 
@@ -147,7 +122,6 @@ func UseRemote(inOutCtx *context.Context, host string, client *http.Client) (err
 			}
 			defer rsp.Body.Close()
 		} else {
-			aeCtx := AEContextNoTxn(*inOutCtx)
 			if aeCtx == nil {
 				aeCtx = context.Background()
 			}
@@ -165,3 +139,40 @@ func UseRemote(inOutCtx *context.Context, host string, client *http.Client) (err
 	*inOutCtx = setupAECtx(*inOutCtx, aeCtx)
 	return nil
 }
+
+// prodState is the current production state.
+type prodState struct {
+	// ctx is the current derived GAE context.
+	ctx context.Context
+
+	// noTxnCtx is a Context maintained alongside ctx. When a transaction is
+	// entered, ctx will be updated, but noTxnCtx will not, allowing extra-
+	// transactional Context access.
+	noTxnCtx context.Context
+
+	// inTxn if true if this is in a transaction, false otherwise.
+	inTxn bool
+}
+
+func getProdState(c context.Context) prodState {
+	if v := c.Value(prodStateKey).(*prodState); v != nil {
+		return *v
+	}
+	return prodState{}
+}
+
+func withProdState(c context.Context, ps prodState) context.Context {
+	return context.WithValue(c, prodStateKey, &ps)
+}
+
+func (ps *prodState) context(c context.Context) context.Context {

[dnj, 2016/09/01 15:25:40]

We *should* probably do cancellation at some point. But that'd require its own
goroutine and junk so unless GAE main library supports it (which, I guess, we'll
never know since we don't offer it the opportunity to support it), this is fine.

[iannucci, 2016/09/16 01:01:13]

Hm I think this method needs a docstring or something. It does say 'context' 5
times in the function declaration, but I'm still not sure what it's supposed to
do :)

[dnj, 2016/09/16 05:44:43]

context context context context context context context context context context
context context context context context context context context

+	aeCtx := ps.ctx
+	if aeCtx == nil {
+		return nil
+	}
+
+	if deadline, ok := c.Deadline(); ok {
+		aeCtx, _ = context.WithDeadline(aeCtx, deadline)
+	}
+	return aeCtx
+}