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..e9fb25697d66befd5b6c34ec10d0363d600b8325 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
-	prodContextNoTxnKey key = 1
-	probeCacheKey       key = 2
+	prodStateKey  = "contains the current *prodState"
+	probeCacheKey = "contains the current *infoProbeCache"
 )
 
 // 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)
+
 		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,45 @@ 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)
+}
+
+// context returns the current AppEngine-bound Context. Prior to returning,
+// the deadline from "c" (if any) is applied.
+//
+// Note that this does not (currently) apply any other Done state or propagate
+// cancellation from "c".

[iannucci, 2016/09/16 07:39:43]

that's unfortunate :/. I think it would be good to address this now (or at least
file a bug/followup CL so we don't forget about it).

It looks like Context.Done() is supposed to return nil if the context cannot be
cancelled. That would at least allow us to only burn a goroutine in the event
that the parent context is even cancellable.

[dnj, 2016/09/16 16:09:21]

This is consistent with pre-patch behavior, so I'm reluctant to add noise from
this change to this already huge patch, especially since such a change would be
API-compatible. Filed: https://github.com/luci/gae/issues/59

+func (ps *prodState) context(c 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
+}