Interface update, per-method Contexts.

filter/txnBuf/state.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 txnBuf
 
 import (
         "bytes"
         "sync"
 
@@ skipping 87 matching lines @@
         sync.Mutex
 
         // encoded key -> size of entity. A size of 0 means that the entity is
         // deleted.
         entState *sizeTracker
         bufDS    datastore.RawInterface
 
         roots     stringset.Set
         rootLimit int
 
-        aid      string
-        ns       string
+        kc       datastore.KeyContext
         parentDS datastore.RawInterface
 
         // sizeBudget is the number of bytes that this transaction has to operate
         // within. It's only used when attempting to apply() the transaction, and
         // it is the threshold for the delta of applying this transaction to the
         // parent transaction. Note that a buffered transaction could actually have
         // a negative delta if the parent transaction had many large entities which
         // the inner transaction deleted.
         sizeBudget int64
         // countBudget is the number of entity writes that this transaction has to
         // operate in.
         writeCountBudget int
 }
 
 func withTxnBuf(ctx context.Context, cb func(context.Context) error, opts *datastore.TransactionOptions) error {
-        inf := info.Get(ctx)
-        ns, _ := inf.GetNamespace()
-
         parentState, _ := ctx.Value(dsTxnBufParent).(*txnBufState)
         roots := stringset.New(0)
         rootLimit := 1
         if opts != nil && opts.XG {
                 rootLimit = XGTransactionGroupLimit
         }
         sizeBudget, writeCountBudget := DefaultSizeBudget, DefaultWriteCountBudget
         if parentState != nil {
                 // TODO(riannucci): this is a bit wonky since it means that a child
                 // transaction declaring XG=true will only get to modify 25 groups IF
                 // they're same groups affected by the parent transactions. So instead of
                 // respecting opts.XG for inner transactions, we just dup everything from
                 // the parent transaction.
                 roots = parentState.roots.Dup()
                 rootLimit = parentState.rootLimit
 
                 sizeBudget = parentState.sizeBudget - parentState.entState.total
                 writeCountBudget = parentState.writeCountBudget - parentState.entState.numWrites()
         }
 
         state := &txnBufState{
                 entState:         &sizeTracker{},
-                bufDS:            memory.NewDatastore(inf).Raw(),
+                bufDS:            memory.NewDatastore(ctx, info.Raw(ctx)),
                 roots:            roots,
                 rootLimit:        rootLimit,
-                ns:               ns,
-                aid:              inf.FullyQualifiedAppID(),
-                parentDS:         datastore.Get(context.WithValue(ctx, dsTxnBufHaveLock, true)).Raw(),
+                kc:               datastore.GetKeyContext(ctx),
+                parentDS:         datastore.Raw(context.WithValue(ctx, dsTxnBufHaveLock, true)),
                 sizeBudget:       sizeBudget,
                 writeCountBudget: writeCountBudget,
         }
         if err := cb(context.WithValue(ctx, dsTxnBufParent, state)); err != nil {
                 return err
         }
 
         // no reason to unlock this ever. At this point it's toast.
         state.Lock()
 
@@ skipping 339 matching lines @@
 
         err = t.bufDS.Run(fq, func(key *datastore.Key, data datastore.PropertyMap, _ datastore.CursorCB) error {
                 toPutKeys = append(toPutKeys, key)
                 toPut = append(toPut, data)
                 return nil
         })
         memoryCorruption(err)
 
         for keyStr, size := range t.entState.keyToSize {
                 if size == 0 {
-                        k, err := serialize.ReadKey(bytes.NewBufferString(keyStr), serialize.WithoutContext, t.aid, t.ns)
+                        k, err := serialize.ReadKey(bytes.NewBufferString(keyStr), serialize.WithoutContext, t.kc)
                         memoryCorruption(err)
                         toDel = append(toDel, k)
                 }
         }
 
         return
 }
 
 func (t *txnBufState) canApplyLocked(s *txnBufState) error {
         proposedState := t.entState.dup()
@@ skipping 38 matching lines @@
 // plus a stringset of all the encoded root keys that `keys` represents.
 func toEncoded(keys []*datastore.Key) (full []string, roots stringset.Set) {
         roots = stringset.New(len(keys))
         full = make([]string, len(keys))
         for i, k := range keys {
                 roots.Add(string(serialize.ToBytes(k.Root())))
                 full[i] = string(serialize.ToBytes(k))
         }
         return
 }