Add transaction buffer filter.

filter/txnBuf/state.go

Index: filter/txnBuf/state.go
diff --git a/filter/txnBuf/state.go b/filter/txnBuf/state.go
new file mode 100644
index 0000000000000000000000000000000000000000..d30eea1a74e77e4737cde10f5f6003b3c01ee7c8
--- /dev/null
+++ b/filter/txnBuf/state.go
@@ -0,0 +1,379 @@
+// 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 txnBuf
+
+import (
+	"bytes"
+	"sync"
+
+	"github.com/luci/gae/impl/memory"
+	"github.com/luci/gae/service/datastore"
+	"github.com/luci/gae/service/datastore/serialize"
+	"github.com/luci/gae/service/info"
+	"github.com/luci/luci-go/common/errors"
+	"github.com/luci/luci-go/common/stringset"
+	"golang.org/x/net/context"
+)
+
+// DefaultSizeBudget is the size budget for the root transaction.
+//
+// Because our estimation algorithm isn't entirely correct, we take 5% off
+// the limit for encoding and estimate inaccuracies.
+//
+// 10MB taken on 2015/09/24:
+// https://cloud.google.com/appengine/docs/go/datastore/#Go_Quotas_and_limits
+const DefaultSizeBudget = int64((10 * 1000 * 1000) * 0.95)
+
+// DefaultSizeThreshold prevents the root transaction from getting too close
+// to the budget. If the code attempts to begin a transaction which would have
+// less than this threshold for its budget, the transaction will immediately
+// return ErrTransactionTooLarge.
+const DefaultSizeThreshold = int64(10 * 1000)
+
+// XGTransactionGroupLimit is the number of transaction groups to allow in an
+// XG transaction.
+//
+// 25 taken on 2015/09/24:
+// https://cloud.google.com/appengine/docs/go/datastore/transactions#Go_What_can_be_done_in_a_transaction
+const XGTransactionGroupLimit = 25
+
+type sizeTracker struct {
+	keyToSize map[string]int64
+	total     int64
+}
+
+func (s *sizeTracker) set(key string, val int64) {
+	prev, existed := s.keyToSize[key]
+	if s.keyToSize == nil {
+		s.keyToSize = make(map[string]int64)
+	}
+	s.keyToSize[key] = val
+	s.total += val - prev
+	if !existed {
+		s.total += int64(len(key))
+	}
+}
+
+func (s *sizeTracker) get(key string) (int64, bool) {
+	size, has := s.keyToSize[key]
+	return size, has
+}
+
+func (s *sizeTracker) has(key string) bool {
+	_, has := s.keyToSize[key]
+	return has
+}
+
+func (s *sizeTracker) dup() *sizeTracker {
+	if len(s.keyToSize) == 0 {
+		return &sizeTracker{}
+	}
+	k2s := make(map[string]int64, len(s.keyToSize))
+	for k, v := range s.keyToSize {
+		k2s[k] = v
+	}
+	return &sizeTracker{k2s, s.total}
+}
+
+type txnBufState struct {
+	sync.Mutex
+
+	// encoded key -> size of entity. A size of 0 means that the entity is
+	// deleted.
+	entState *sizeTracker
+	memDS    datastore.RawInterface
+
+	roots     stringset.Set
+	rootLimit int
+
+	aid         string
+	ns          string
+	parentDS    datastore.RawInterface
+	parentState *txnBufState
+
+	// 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
+
+	siblingLock sync.Mutex
+}
+
+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 := DefaultSizeBudget
+	if parentState != nil {
+		parentState.siblingLock.Lock()
+		defer parentState.siblingLock.Unlock()
+
+		// 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
+		if sizeBudget < DefaultSizeThreshold {
+			return ErrTransactionTooLarge
+		}
+	}
+
+	memDS := memory.NewDatastore(ns)
+	{

[Vadim Sh., 2015/09/28 18:52:56]

Remove...

[iannucci, 2015/09/29 03:21:38]

k

+		t := memDS.Testable()
+		t.Consistent(true)
+		t.AutoIndex(true)
+		t.DisableSpecialEntities(true)
+	}
+	state := &txnBufState{
+		entState:    &sizeTracker{},
+		memDS:       memDS,
+		roots:       roots,
+		rootLimit:   rootLimit,
+		ns:          ns,
+		aid:         inf.AppID(),
+		parentDS:    datastore.Get(ctx).Raw(),
+		parentState: parentState,
+		sizeBudget:  sizeBudget,
+	}
+	ctx = context.WithValue(ctx, dsTxnBufParent, state)
+	err := cb(ctx)
+	if err != nil {
+		return err
+	}
+	return state.apply()
+}
+
+type item struct {
+	key      *datastore.Key
+	data     datastore.PropertyMap
+	buffered bool
+
+	encKey string
+	cmpRow string
+	err    error
+}
+
+func (i *item) getEncKey() string {
+	if i.encKey == "" {
+		i.encKey = string(serialize.ToBytes(i.key))
+	}
+	return i.encKey
+}
+
+func (i *item) getCmpRow(lower, upper []byte, order []datastore.IndexColumn) string {
+	if i.cmpRow == "" {
+		row, key := toComparableString(lower, upper, order, i.key, i.data)
+		i.cmpRow = string(row)
+		if i.encKey == "" {
+			i.encKey = string(key)
+		}
+	}
+	return i.cmpRow
+}
+
+func (t *txnBufState) updateRootsLocked(roots stringset.Set) error {
+	curRootLen := t.roots.Len()
+	proposedRoots := stringset.New(1)
+	roots.Iter(func(root string) bool {
+		if !t.roots.Has(root) {
+			proposedRoots.Add(root)
+		}
+		return proposedRoots.Len()+curRootLen <= t.rootLimit
+	})
+	if proposedRoots.Len()+curRootLen > t.rootLimit {
+		return errors.New("operating on too many entity groups in nested transaction")
+	}
+	// only need to update the roots if they did something that required updating
+	if proposedRoots.Len() > 0 {
+		proposedRoots.Iter(func(root string) bool {
+			t.roots.Add(root)
+			return true
+		})
+	}
+	return nil
+}
+
+func (t *txnBufState) getMulti(keys []*datastore.Key) ([]item, error) {
+	encKeys, roots := toEncoded(keys)
+	ret := make([]item, len(keys))
+
+	idxMap := []int(nil)
+	toGetKeys := []*datastore.Key(nil)
+
+	t.Lock()
+	defer t.Unlock()
+
+	if err := t.updateRootsLocked(roots); err != nil {
+		return nil, err
+	}
+
+	for i, key := range keys {
+		ret[i].key = key
+		ret[i].encKey = encKeys[i]
+		if size, ok := t.entState.get(ret[i].getEncKey()); ok {
+			ret[i].buffered = true
+			if size > 0 {
+				idxMap = append(idxMap, i)
+				toGetKeys = append(toGetKeys, key)
+			}
+		}
+	}
+
+	if len(toGetKeys) > 0 {
+		j := 0
+		t.memDS.GetMulti(toGetKeys, nil, func(pm datastore.PropertyMap, err error) {
+			idx := idxMap[j]
+			impossible(err)

[Vadim Sh., 2015/09/28 18:52:56]

nit: move it up one line

[iannucci, 2015/09/29 03:21:38]

done

+			ret[idx].data = pm
+			j++
+		})
+	}
+
+	return ret, nil
+}
+
+func (t *txnBufState) deleteMulti(keys []*datastore.Key) error {
+	encKeys, roots := toEncoded(keys)
+
+	t.Lock()
+	defer t.Unlock()
+
+	if err := t.updateRootsLocked(roots); err != nil {
+		return err
+	}
+
+	i := 0
+	err := t.memDS.DeleteMulti(keys, func(err error) {
+		impossible(err)
+		t.entState.set(encKeys[i], 0)
+		i++
+	})
+	impossible(err)
+	return nil
+}
+
+func (t *txnBufState) putMulti(keys []*datastore.Key, vals []datastore.PropertyMap) error {
+	encKeys, roots := toEncoded(keys)
+
+	t.Lock()
+	defer t.Unlock()
+
+	if err := t.updateRootsLocked(roots); err != nil {
+		return err
+	}
+
+	i := 0
+	err := t.memDS.PutMulti(keys, vals, func(k *datastore.Key, err error) {
+		impossible(err)
+		t.entState.set(encKeys[i], vals[i].EstimateSize())
+		i++
+	})
+	impossible(err)
+	return nil
+}
+
+func (t *txnBufState) apply() error {
+	t.Lock()
+	defer t.Unlock()
+
+	// if parentState is nil... just try to commit this anyway. The estimates
+	// we're using here are just educated guesses. If it fits for real, then
+	// hooray. If not, then the underlying datastore will error.
+	if t.parentState != nil {
+		proposedState := t.parentState.entState.dup()
+		for k, v := range t.entState.keyToSize {
+			proposedState.set(k, v)
+		}
+		if proposedState.total > t.sizeBudget {
+			return ErrTransactionTooLarge
+		}
+	}
+
+	toPutKeys := []*datastore.Key(nil)
+	toPut := []datastore.PropertyMap(nil)
+	toDel := []*datastore.Key(nil)
+
+	fq, err := datastore.NewQuery("").Finalize()

[Vadim Sh., 2015/09/28 18:52:56]

wat?

[iannucci, 2015/09/29 03:21:38]

added comment

+	impossible(err)
+
+	err = t.memDS.Run(fq, func(key *datastore.Key, data datastore.PropertyMap, _ datastore.CursorCB) bool {
+		toPutKeys = append(toPutKeys, key)
+		toPut = append(toPut, data)
+		return true
+	})
+	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)
+			memoryCorruption(err)
+			toDel = append(toDel, k)
+		}
+	}
+
+	wg := sync.WaitGroup{}
+
+	pErr := error(nil)
+	dErr := error(nil)
+
+	ds := t.parentDS
+	if toPut != nil {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			mErr := errors.NewLazyMultiError(len(toPut))
+			i := 0
+			pErr = ds.PutMulti(toPutKeys, toPut, func(_ *datastore.Key, err error) {
+				i++
+				mErr.Assign(i, err)
+			})
+			pErr = mErr.Get()
+		}()
+	}
+
+	if toDel != nil {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			mErr := errors.NewLazyMultiError(len(toDel))
+			i := 0
+			dErr = ds.DeleteMulti(toDel, func(err error) {
+				mErr.Assign(i, err)
+				i++
+			})
+			dErr = mErr.Get()
+		}()
+	}
+	wg.Wait()
+
+	if pErr != nil {
+		return pErr
+	}
+	return dErr
+}
+
+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
+}