Interface update, per-method Contexts.

impl/memory/datastore_data.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 (
         "bytes"
         "fmt"
         "strings"
         "sync"
-        "sync/atomic"
 
         ds "github.com/luci/gae/service/datastore"
         "github.com/luci/gae/service/datastore/serialize"
         "github.com/luci/luci-go/common/errors"
+
         "golang.org/x/net/context"
 )
 
 //////////////////////////////// dataStoreData /////////////////////////////////
 
 type dataStoreData struct {
         rwlock sync.RWMutex
 
         // the 'appid' of this datastore
         aid string
@@ skipping 142 matching lines @@
 func (d *dataStoreData) namespaces() []string {
         d.rwlock.Lock()
         defer d.rwlock.Unlock()
 
         return namespaces(d.head)
 }
 
 /////////////////////////// indexes(dataStoreData) ////////////////////////////
 
 func groupMetaKey(key *ds.Key) []byte {
-        return keyBytes(ds.NewKey("", "", "__entity_group__", "", 1, key.Root()))
+        return keyBytes(ds.KeyContext{}.NewKey("__entity_group__", "", 1, key.Root()))
 }
 
 func groupIDsKey(key *ds.Key) []byte {
-        return keyBytes(ds.NewKey("", "", "__entity_group_ids__", "", 1, key.Root()))
+        return keyBytes(ds.KeyContext{}.NewKey("__entity_group_ids__", "", 1, key.Root()))
 }
 
 func rootIDsKey(kind string) []byte {
-        return keyBytes(ds.NewKey("", "", "__entity_root_ids__", kind, 0, nil))
+        return keyBytes(ds.KeyContext{}.NewKey("__entity_root_ids__", kind, 0, nil))
 }
 
 func curVersion(ents memCollection, key []byte) int64 {
         if ents != nil {
                 if v := ents.Get(key); v != nil {
                         pm, err := rpm(v)
                         memoryCorruption(err)
 
                         pl, ok := pm["__version__"]
                         if ok && len(pl) > 0 && pl[0].Type() == ds.PTInt {
@@ skipping 74 matching lines @@
         }
         return incrementLocked(ents, idKey, n), nil
 }
 
 func (d *dataStoreData) fixKeyLocked(ents memCollection, key *ds.Key) (*ds.Key, error) {
         if key.IsIncomplete() {
                 id, err := d.allocateIDsLocked(ents, key, 1)
                 if err != nil {
                         return key, err
                 }
-                key = ds.NewKey(key.AppID(), key.Namespace(), key.Kind(), "", id, key.Parent())
+                key = key.KeyContext().NewKey(key.Kind(), "", id, key.Parent())
         }
         return key, nil
 }
 
 func (d *dataStoreData) putMulti(keys []*ds.Key, vals []ds.PropertyMap, cb ds.NewKeyCB) error {
         ns := keys[0].Namespace()
 
         for i, k := range keys {
                 pmap, _ := vals[i].Save(false)
                 dataBytes := serialize.ToBytes(pmap)
@@ skipping 112 matching lines @@
 }
 
 func (d *dataStoreData) canApplyTxn(obj memContextObj) bool {
         // TODO(riannucci): implement with Flush/FlushRevert for persistance.
 
         txn := obj.(*txnDataStoreData)
         for rk, muts := range txn.muts {
                 if len(muts) == 0 { // read-only
                         continue
                 }
-                prop, err := serialize.ReadProperty(bytes.NewBufferString(rk), serialize.WithContext, "", "")
+                prop, err := serialize.ReadProperty(bytes.NewBufferString(rk), serialize.WithContext, ds.KeyContext{})
                 memoryCorruption(err)
 
                 k := prop.Value().(*ds.Key)
 
                 entKey := "ents:" + k.Namespace()
                 mkey := groupMetaKey(k)
                 entsHead := d.head.GetCollection(entKey)
                 entsSnap := txn.snap.GetCollection(entKey)
                 vHead := curVersion(entsHead, mkey)
                 vSnap := curVersion(entsSnap, mkey)
@@ skipping 22 matching lines @@
                         }
                 }
         }
 }
 
 func (d *dataStoreData) mkTxn(o *ds.TransactionOptions) memContextObj {
         return &txnDataStoreData{
                 // alias to the main datastore's so that testing code can have primitive
                 // access to break features inside of transactions.
                 parent: d,
-                isXG:   o != nil && o.XG,
-                snap:   d.head.Snapshot(),
-                muts:   map[string][]txnMutation{},
+                txn: &transactionImpl{
+                        isXG: o != nil && o.XG,
+                },
+                snap: d.head.Snapshot(),
+                muts: map[string][]txnMutation{},
         }
 }
 
 func (d *dataStoreData) endTxn() {}
 
 /////////////////////////////// txnDataStoreData ///////////////////////////////
 
 type txnMutation struct {
         key  *ds.Key
         data ds.PropertyMap
 }
 
 type txnDataStoreData struct {
         sync.Mutex
 
         parent *dataStoreData
 
-        // boolean 0 or 1, use atomic.*Int32 to access.
-        closed int32
-        isXG   bool
+        txn *transactionImpl
 
         snap memStore
 
         // string is the raw-bytes encoding of the entity root incl. namespace
         muts map[string][]txnMutation
         // TODO(riannucci): account for 'transaction size' limit of 10MB by summing
         // length of encoded keys + values.
 }
 
 var _ memContextObj = (*txnDataStoreData)(nil)
 
 const xgEGLimit = 25
 
 func (*txnDataStoreData) canApplyTxn(memContextObj) bool { return false }
+
 func (td *txnDataStoreData) endTxn() {
-        if atomic.LoadInt32(&td.closed) == 1 {
-                panic("cannot end transaction twice")
+        if err := td.txn.close(); err != nil {
+                panic(err)
         }
-        atomic.StoreInt32(&td.closed, 1)
 }
 func (*txnDataStoreData) applyTxn(context.Context, memContextObj) {
         impossible(fmt.Errorf("cannot create a recursive transaction"))
 }
 func (*txnDataStoreData) mkTxn(*ds.TransactionOptions) memContextObj {
         impossible(fmt.Errorf("cannot create a recursive transaction"))
         return nil
 }
 
 func (td *txnDataStoreData) run(f func() error) error {
         // Slightly different from the SDK... datastore and taskqueue each implement
         // this here, where in the SDK only datastore.transaction.Call does.
-        if atomic.LoadInt32(&td.closed) == 1 {
-                return errors.New("datastore: transaction context has expired")
+        if err := td.txn.valid(); err != nil {
+                return err
         }
         return f()
 }
 
 // writeMutation ensures that this transaction can support the given key/value
 // mutation.
 //
 //   if getOnly is true, don't record the actual mutation data, just ensure that
 //         the key is in an included entity group (or add an empty entry for that
 //         group).
 //
 //   if !getOnly && data == nil, this counts as a deletion instead of a Put.
 //
 // Returns an error if this key causes the transaction to cross too many entity
 // groups.
 func (td *txnDataStoreData) writeMutation(getOnly bool, key *ds.Key, data ds.PropertyMap) error {
         rk := string(keyBytes(key.Root()))
 
         td.Lock()
         defer td.Unlock()
 
         if _, ok := td.muts[rk]; !ok {
                 limit := 1
-                if td.isXG {
+                if td.txn.isXG {
                         limit = xgEGLimit
                 }
                 if len(td.muts)+1 > limit {
                         msg := "cross-group transaction need to be explicitly specified (xg=True)"
-                        if td.isXG {
+                        if td.txn.isXG {
                                 msg = "operating on too many entity groups in a single transaction"
                         }
                         return errors.New(msg)
                 }
                 td.muts[rk] = []txnMutation{}
         }
         if !getOnly {
                 td.muts[rk] = append(td.muts[rk], txnMutation{key, data})
         }
 
@@ skipping 43 matching lines @@
         }
         return nil
 }
 
 func keyBytes(key *ds.Key) []byte {
         return serialize.ToBytes(ds.MkProperty(key))
 }
 
 func rpm(data []byte) (ds.PropertyMap, error) {
         return serialize.ReadPropertyMap(bytes.NewBuffer(data),
-                serialize.WithContext, "", "")
+                serialize.WithContext, ds.KeyContext{})
 }
 
 func namespaces(store memStore) []string {
         var namespaces []string
         for _, c := range store.GetCollectionNames() {
                 ns, has := trimPrefix(c, "ents:")
                 if !has {
                         if len(namespaces) > 0 {
                                 break
                         }
                         continue
                 }
                 namespaces = append(namespaces, ns)
         }
         return namespaces
 }
 
 func trimPrefix(v, p string) (string, bool) {
         if strings.HasPrefix(v, p) {
                 return v[len(p):], true
         }
         return v, false
 }