Simple memory testing for gae/wrapper

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

Index: go/src/infra/gae/libs/wrapper/memory/context.go
diff --git a/go/src/infra/gae/libs/wrapper/memory/context.go b/go/src/infra/gae/libs/wrapper/memory/context.go
new file mode 100644
index 0000000000000000000000000000000000000000..17cf5ea221e3acf090128aa2ac88c02f731a184b
--- /dev/null
+++ b/go/src/infra/gae/libs/wrapper/memory/context.go
@@ -0,0 +1,162 @@
+// 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 (
+	"fmt"
+	"math/rand"
+	"sync"
+
+	"golang.org/x/net/context"
+
+	"appengine/datastore"
+
+	"infra/gae/libs/wrapper"
+)
+
+type memContextObj interface {
+	sync.Locker
+	canApplyTxn(m memContextObj) bool
+	applyTxn(rnd *rand.Rand, m memContextObj)
+
+	endTxn()
+	mkTxn(*datastore.TransactionOptions) (memContextObj, error)
+}
+
+type memContext []memContextObj
+
+func newMemContext() memContext {
+	return memContext{
+		newTaskQueueData(),
+		newDataStoreData(),
+	}
+}
+
+var memContextIndices = map[string]int{
+	"TQ": 0,
+	"DS": 1,
+}
+
+func (m memContext) Get(item string) memContextObj {
+	if i, ok := memContextIndices[item]; !ok {
+		panic(fmt.Errorf("wrapper/memory: cannot get context item %q", item))

[M-A Ruel, 2015/05/25 18:21:08]

Why not return nil?

[iannucci, 2015/05/27 19:33:30]

Because it's meaningless. This object is only ever constructed with 2 objects in
it (for now, at some point later it will have 3), and those two objects go in
specific slots. If we can't find the item, it's an irrecoverable
program-terminating error. That's what panic is for.

+	} else {
+		return m[i]
+	}
+}
+
+func (m memContext) Lock() {
+	for _, itm := range m {
+		itm.Lock()
+	}
+}
+
+func (m memContext) Unlock() {
+	for i := len(m) - 1; i >= 0; i-- {
+		m[i].Unlock()
+	}
+}
+
+func (m memContext) endTxn() {
+	for _, itm := range m {
+		itm.endTxn()
+	}
+}
+
+func (m memContext) mkTxn(o *datastore.TransactionOptions) (ret memContext, err error) {
+	for _, itm := range m {
+		newItm, err := itm.mkTxn(o)

[M-A Ruel, 2015/05/25 18:21:08]

You alias err here. I'd prefer to not use named return argument and instead
insert at the top of the function:

ret := make(memContext, 0, len(m))

it will reduce memory pressure in the success control flow.

[iannucci, 2015/05/27 19:33:30]

done, ditched append too since we don't really need it.

+		if err != nil {
+			return nil, err
+		}
+		ret = append(ret, newItm)
+	}
+	return ret, nil
+}
+
+func (m memContext) canApplyTxn(txnCtx memContext) bool {
+	for i := range m {

[M-A Ruel, 2015/05/25 18:21:08]

Why happens if len(txnCtr) < len(m), panic. :)

[iannucci, 2015/05/27 19:33:30]

yes, that's intentional.

+		if !m[i].canApplyTxn(txnCtx[i]) {
+			return false
+		}
+	}
+	return true
+}
+
+func (m memContext) applyTxn(rnd *rand.Rand, txnCtx memContext) {
+	for i := range m {

[M-A Ruel, 2015/05/25 18:21:08]

same 

[iannucci, 2015/05/27 19:33:30]

still intentional :)

+		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(
+		context.WithValue(c, memContextKey, newMemContext()),
+		giContextKey, newGlobalInfoData())
+}
+
+// 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))))
+}
+
+func cur(c context.Context) (p memContext) {
+	p, _ = c.Value(memContextKey).(memContext)
+	return
+}
+
+type memContextKeyType int
+
+var memContextKey memContextKeyType
+
+// weird stuff
+
+// RunInTransaction is here because it's really a service-wide transaction, not
+// just in the datastore. TaskQueue behaves differently in a transaction in
+// a couple ways, for example.
+//
+// It really should have been appengine.Context.RunInTransaction(func(tc...)),
+// but because it's not, this method is on dsImpl instead to mirror the official
+// API.
+//
+// The fake implementation also differs from the real implementation because the
+// fake TaskQueue is NOT backed by the fake Datastore. This is done to make the
+// test-access API for TaskQueue better (instead of trying to reconstitute the
+// state of the task queue from a bunch of datastore accesses).
+func (d *dsImpl) RunInTransaction(f func(context.Context) error, o *datastore.TransactionOptions) (err error) {
+	curMC := cur(d.c)
+
+	txnMC, err := curMC.mkTxn(o)

[M-A Ruel, 2015/05/25 18:21:08]

You alias err here, so don't name the return variable.

[iannucci, 2015/05/27 19:33:30]

no I don't (http://play.golang.org/p/ZQcAvmI7WQ), but I also don't rely on the
named variable, so I removed it.

+	if err != nil {
+		return err
+	}
+
+	defer func() {
+		txnMC.Lock()
+		defer txnMC.Unlock()
+
+		txnMC.endTxn()
+	}()
+
+	err = f(context.WithValue(d.c, memContextKey, txnMC))
+	if err != nil {

[M-A Ruel, 2015/05/25 18:21:08]

if err = f(context.WithValue(d.c, memContextKey, txnMC)); err != nil {

[iannucci, 2015/05/27 19:33:30]

done

+		return err
+	}
+
+	return func() error {

[M-A Ruel, 2015/05/25 18:21:08]

I don't see why.

[iannucci, 2015/05/27 19:33:30]

used to have code after the function and needed to bound the lock. removed.

+		txnMC.Lock()
+		defer txnMC.Unlock()
+
+		if curMC.canApplyTxn(txnMC) {
+			curMC.applyTxn(wrapper.GetMathRand(d.c), txnMC)
+		} else {
+			return datastore.ErrConcurrentTransaction
+		}
+		return nil
+	}()
+}