Simplify memory package interface.

go/src/infra/gae/libs/wrapper/memory/context.go

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package memory
 
 import (
+        "errors"
         "infra/gae/libs/wrapper"
         "math/rand"
         "sync"
 
         "golang.org/x/net/context"
 
         "appengine/datastore"
 )
 
 type memContextObj interface {
@@ skipping 67 matching lines @@
         return true
 }
 
 func (m memContext) applyTxn(rnd *rand.Rand, txnCtxObj memContextObj) {
         txnCtx := txnCtxObj.(memContext)
         for i := range m {
                 m[i].applyTxn(rnd, txnCtx[i])
         }
 }
 
-// Enable adds a new memory context to c. This new memory context will have
-// a zeroed state.
-func Enable(c context.Context) context.Context {
-        return context.WithValue(
+// Use adds implementations for the following gae/wrapper interfaces to the
+// context:
+//   * wrapper.Datastore
+//   * wrapper.TaskQueue
+//   * wrapper.Memcache
+//   * wrapper.GlobalInfo
+//
+// These can be retrieved with the "gae/wrapper".Get functions.
+//
+// The implementations are all backed by an in-memory implementation, and start
+// with an empty state.
+//
+// Using this more than once per context.Context will cause a panic.
+func Use(c context.Context) context.Context {
+        if c.Value(memContextKey) != nil {
+                panic(errors.New("memory.Use: called twice on the same Context"))
+        }
+        c = context.WithValue(
                 context.WithValue(c, memContextKey, newMemContext()),
                 giContextKey, &globalInfoData{})
-}
-
-// Use calls ALL of this packages Use* methods on c. This enables all
-// gae/wrapper Get* api's.
-func Use(c context.Context) context.Context {
-        return UseTQ(UseDS(UseMC(UseGI(c))))
+        return useTQ(useDS(useMC(useGI(c))))
 }
 
 func cur(c context.Context) (p memContext) {
         p, _ = c.Value(memContextKey).(memContext)
         return
 }
 
 type memContextKeyType int
 
 var memContextKey memContextKeyType
@@ skipping 34 matching lines @@
         txnMC.Lock()
         defer txnMC.Unlock()
 
         if curMC.canApplyTxn(txnMC) {
                 curMC.applyTxn(wrapper.GetMathRand(d.c), txnMC)
         } else {
                 return datastore.ErrConcurrentTransaction
         }
         return nil
 }