Interface update, per-method Contexts.

impl/prod/info.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 prod
 
 import (
         "time"
 
         "github.com/luci/gae/service/info"
@@ skipping 22 matching lines @@
 
 func (g giImpl) AccessToken(scopes ...string) (token string, expiry time.Time, err error) {
         return appengine.AccessToken(g.aeCtx, scopes...)
 }
 func (g giImpl) AppID() string {
         return appengine.AppID(g.aeCtx)
 }
 func (g giImpl) FullyQualifiedAppID() string {
         return getProbeCache(g.usrCtx).fqaid
 }
-func (g giImpl) GetNamespace() (string, bool) {
-        pc := getProbeCache(g.usrCtx)
-        if ns := pc.namespace; ns != nil {
-                return *ns, true
-        }
-        return "", false
+func (g giImpl) GetNamespace() string {
+        return getProbeCache(g.usrCtx).namespace
 }
 func (g giImpl) Datacenter() string {
         return appengine.Datacenter(g.aeCtx)
 }
 func (g giImpl) DefaultVersionHostname() string {
         return appengine.DefaultVersionHostname(g.aeCtx)
 }
 func (g giImpl) InstanceID() string {
         return appengine.InstanceID()
 }
 func (g giImpl) IsDevAppServer() bool {
         return appengine.IsDevAppServer()
 }
 func (g giImpl) IsOverQuota(err error) bool {
         return appengine.IsOverQuota(err)
 }
 func (g giImpl) IsTimeoutError(err error) bool {
         return appengine.IsTimeoutError(err)
 }
 func (g giImpl) ModuleHostname(module, version, instance string) (string, error) {
         return appengine.ModuleHostname(g.aeCtx, module, version, instance)
 }
 func (g giImpl) ModuleName() (name string) {
         return appengine.ModuleName(g.aeCtx)
 }
 func (g giImpl) Namespace(namespace string) (context.Context, error) {
-        aeCtx, err := appengine.Namespace(g.aeCtx, namespace)
+        c := g.usrCtx
+
+        pc := *getProbeCache(c)
+        if pc.namespace == namespace {
+                // Already using this namespace.
+                return c, nil
+        }
+        pc.namespace = namespace
+
+        // Apply the namespace to our retained GAE Contexts.
+        var err error
+        ps := getProdState(c)
+
+        // Apply to current GAE Context.
+        if ps.ctx, err = appengine.Namespace(ps.ctx, namespace); err != nil {
+                return c, err
+        }
+
+        // Apply to non-transactional Context. Since the previous one applied
+        // successfully, this must succeed.
+        ps.noTxnCtx, err = appengine.Namespace(ps.noTxnCtx, namespace)
         if err != nil {
-                return g.usrCtx, err
+                panic(err)
         }
-        usrCtx := context.WithValue(g.usrCtx, prodContextKey, aeCtx)
-        pc := *getProbeCache(usrCtx)
-        pc.namespace = &namespace
-        return withProbeCache(usrCtx, &pc), nil
+
+        // Update our user Context with the new namespace-imbued objects.
+        c = withProbeCache(c, &pc)
+        c = withProdState(c, ps)
+        return c, nil
 }
 func (g giImpl) PublicCertificates() ([]info.Certificate, error) {
         certs, err := appengine.PublicCertificates(g.aeCtx)
         if err != nil {
                 return nil, err
         }
         ret := make([]info.Certificate, len(certs))
         for i, c := range certs {
                 ret[i] = info.Certificate(c)
         }
@@ skipping 13 matching lines @@
         }
         return appengine.ServiceAccount(g.aeCtx)
 }
 func (g giImpl) SignBytes(bytes []byte) (keyName string, signature []byte, err error) {
         return appengine.SignBytes(g.aeCtx, bytes)
 }
 func (g giImpl) VersionID() string {
         return appengine.VersionID(g.aeCtx)
 }
 
-func (g giImpl) Testable() info.Testable {
-        return nil
-}
+func (g giImpl) GetTestable() info.Testable { return nil }
 
 type infoProbeCache struct {
-        namespace *string
+        namespace string
         fqaid     string
 }
 
 func probe(aeCtx context.Context) *infoProbeCache {
         probeKey := datastore.NewKey(aeCtx, "Kind", "id", 0, nil)
-        namespace := probeKey.Namespace()
         ipb := infoProbeCache{
-                fqaid: probeKey.AppID(),
-        }
-        if namespace != "" {
-                ipb.namespace = &namespace
+                fqaid:     probeKey.AppID(),
+                namespace: probeKey.Namespace(),
         }
         return &ipb
 }
 
 func getProbeCache(c context.Context) *infoProbeCache {
-        if pc, ok := c.Value(probeCacheKey).(*infoProbeCache); ok {
+        if pc, ok := c.Value(&probeCacheKey).(*infoProbeCache); ok {
                 return pc
         }
         return nil
 }
 
 func withProbeCache(c context.Context, pc *infoProbeCache) context.Context {
-        return context.WithValue(c, probeCacheKey, pc)
+        return context.WithValue(c, &probeCacheKey, pc)
 }