txnbuf: Add write count budget, size budget.

filter/txnBuf/state.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 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/parallel"
         "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)
+// DefaultWriteCountBudget is the maximum number of entities that can be written
+// in a single call.
+//
+// This is not known to be documented, and has instead been extracted from a
+// datastore error message.
+const DefaultWriteCountBudget = 500
 
 // 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
 
 // sizeTracker tracks the size of a buffered transaction. The rules are simple:
 //   * deletes count for the size of their key, but 0 data
@@ skipping 24 matching lines @@
         size, has := s.keyToSize[key]
         return size, has
 }
 
 // has returns true iff key has a tracked value.
 func (s *sizeTracker) has(key string) bool {
         _, has := s.keyToSize[key]
         return has
 }
 
+// numWrites returns the number of tracked write operations.
+func (s *sizeTracker) numWrites() int {
+        return len(s.keyToSize)
+}
+
 // dup returns a duplicate sizeTracker.
 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}
@@ skipping 14 matching lines @@
         ns       string
         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 := DefaultSizeBudget
+        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
-                if sizeBudget < DefaultSizeThreshold {
-                        return ErrTransactionTooLarge
-                }
+                writeCountBudget = parentState.writeCountBudget - parentState.entState.numWrites()
         }
 
         bufDS, err := memory.NewDatastore(inf.FullyQualifiedAppID(), ns)
         if err != nil {
                 return err
         }
 
         state := &txnBufState{
-                entState:   &sizeTracker{},
-                bufDS:      bufDS.Raw(),
-                roots:      roots,
-                rootLimit:  rootLimit,
-                ns:         ns,
-                aid:        inf.AppID(),
-                parentDS:   datastore.Get(context.WithValue(ctx, dsTxnBufHaveLock, true)).Raw(),
-                sizeBudget: sizeBudget,
+                entState:         &sizeTracker{},
+                bufDS:            bufDS.Raw(),
+                roots:            roots,
+                rootLimit:        rootLimit,
+                ns:               ns,
+                aid:              inf.AppID(),
+                parentDS:         datastore.Get(context.WithValue(ctx, dsTxnBufHaveLock, true)).Raw(),
+                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()
 
         if parentState == nil {
                 return commitToReal(state)
@@ skipping 320 matching lines @@
 
         return
 }
 
 func (t *txnBufState) canApplyLocked(s *txnBufState) error {
         proposedState := t.entState.dup()
 
         for k, v := range s.entState.keyToSize {
                 proposedState.set(k, v)
         }
-        if proposedState.total > s.sizeBudget {
+        switch {
+        case proposedState.numWrites() > t.writeCountBudget:
+                // The new net number of writes must be below the parent's write count
+                // cutoff.
+                fallthrough
+
+        case proposedState.total > t.sizeBudget:
+                // Make sure our new calculated size is within the parent's size budget.
+                //
+                // We have:
+                // - proposedState.total: The "new world" total bytes were this child
+                //   transaction committed to the parent.
+                // - t.sizeBudget: The maximum number of bytes that this parent can
+                //   accommodate.
                 return ErrTransactionTooLarge
         }
+
         return nil
 }
 
 func (t *txnBufState) commitLocked(s *txnBufState) {
         toPut, toPutKeys, toDel := s.effect()
 
         if len(toPut) > 0 {
                 impossible(t.putMulti(toPutKeys, toPut,
                         func(_ *datastore.Key, err error) error { return err }, true))
         }
 
         if len(toDel) > 0 {
                 impossible(t.deleteMulti(toDel, func(err error) error { return err }, true))
         }
 }
 
 // toEncoded returns a list of all of the serialized versions of these keys,
 // 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
 }