storage: add limits to skipped data iteration

Previously when iterating engine using  MVCCIncrementalIterator caller
could skip large amounts of non-matching data which would result in
"unbounded" resource usage.
This is becoming a problem for resource constrained clusters where low
priority requests like export that are used by backups to interfere with
high priority workloads. If we want to throttle backups we need to be able
to limit how many underlying operations we want to perform per request.
This change adds an optional iteration limit to the iterator. Once the
limit is reached, iterator will return an error. Error will provide a
resume key to continue iteration in next request.

Release note: None

pkg/storage/bench_test.go

@@ -555,10 +555,11 @@ func loadTestData(dir string, numKeys, numBatches, batchTimeSpan, valueBytes int
 
 	var batch Batch
 	var minWallTime int64
+	batchSize := len(keys) / numBatches
 	for i, key := range keys {
-		if scaled := len(keys) / numBatches; (i % scaled) == 0 {
+		if (i % batchSize) == 0 {
 			if i > 0 {
-				log.Infof(ctx, "committing (%d/~%d)", i/scaled, numBatches)
+				log.Infof(ctx, "committing (%d/~%d)", i/batchSize, numBatches)
 				if err := batch.Commit(false /* sync */); err != nil {
 					return nil, err
 				}
@@ -568,7 +569,7 @@ func loadTestData(dir string, numKeys, numBatches, batchTimeSpan, valueBytes int
 				}
 			}
 			batch = eng.NewBatch()
-			minWallTime = sstTimestamps[i/scaled]
+			minWallTime = sstTimestamps[i/batchSize]
 		}
 		timestamp := hlc.Timestamp{WallTime: minWallTime + rand.Int63n(int64(batchTimeSpan))}
 		value := roachpb.MakeValueFromBytes(randutil.RandBytes(rng, valueBytes))

pkg/storage/engine.go

@@ -387,6 +387,12 @@ type ExportOptions struct {
 	// would stop immediately when targetSize is reached and return a next versions
 	// timestamp in resumeTs so that subsequent operation can pass it to firstKeyTs.
 	StopMidKey bool
+	// MaxAllowerdIterations limits number of iterations that storage would perform
+	// while colllecting matching values to export. This limit could only be used
+	// together with StopMidKey as it could break iteration over data on any version
+	// and even on a key and timestamp that would not be included in the resulting
+	// export.
+	MaxAllowerdIterations int64
 	// If UseTBI is true, the backing MVCCIncrementalIterator will initialize a
 	// time-bound iterator along with its regular iterator. The TBI will be used
 	// as an optimization to skip over swaths of uninteresting keys i.e. keys

pkg/storage/mvcc_incremental_iterator.go

@@ -11,6 +11,8 @@
 package storage
 
 import (
+	"fmt"
+
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -29,6 +31,8 @@ import (
 //   1. An inline value (non-user data)
 //   2. An intent whose timestamp lies within the time bounds
 //      (if not using enableWriteIntentAggregation)
+//   3. Iterator hits a limit on number of Next calls on underlying
+//      iterator (if limit was requested).
 //
 // Note: The endTime is inclusive to be consistent with the non-incremental
 // iterator, where reads at a given timestamp return writes at that
@@ -93,6 +97,10 @@ type MVCCIncrementalIterator struct {
 	enableWriteIntentAggregation bool
 	// Optional collection of intents created on demand when first intent encountered.
 	intents []roachpb.Intent
+
+	// Underlying iteration limiting.
+	maxAllowedIterations int64
+	iteratorAdvanceCount int64
 }
 
 var _ SimpleMVCCIterator = &MVCCIncrementalIterator{}
@@ -113,6 +121,13 @@ type MVCCIncrementalIterOptions struct {
 	// would be free to decide if it wants to keep collecting entries and intents or
 	// skip entries.
 	EnableWriteIntentAggregation bool
+	// If non zero iterator would only advance or seek up to a certain number of times
+	// before bailing out with ResourceLimitError. When using the feature it is
+	// important to use resume key from the error to ensure forward progress could
+	// be made as advancing up to the first key could already exhaust the limit.
+	// Resume key is not necessarily a valid iteration key as we could stop in between
+	// eligible keys.
+	MaxAllowedIterations int64
 }
 
 // NewMVCCIncrementalIterator creates an MVCCIncrementalIterator with the
@@ -153,6 +168,7 @@ func NewMVCCIncrementalIterator(
 		endTime:                      opts.EndTime,
 		timeBoundIter:                timeBoundIter,
 		enableWriteIntentAggregation: opts.EnableWriteIntentAggregation,
+		maxAllowedIterations:         opts.MaxAllowedIterations,
 	}
 }
 
@@ -176,9 +192,8 @@ func (i *MVCCIncrementalIterator) SeekGE(startKey MVCCKey) {
 		}
 	}
 	i.iter.SeekGE(startKey)
-	if ok, err := i.iter.Valid(); !ok {
-		i.err = err
-		i.valid = false
+	i.iteratorAdvanceCount++
+	if !i.checkValidAndSaveErr() {
 		return
 	}
 	i.err = nil
@@ -199,6 +214,7 @@ func (i *MVCCIncrementalIterator) Close() {
 // key.
 func (i *MVCCIncrementalIterator) Next() {
 	i.iter.Next()
+	i.iteratorAdvanceCount++
 	if !i.checkValidAndSaveErr() {
 		return
 	}
@@ -222,6 +238,7 @@ func (i *MVCCIncrementalIterator) checkValidAndSaveErr() bool {
 // key if the iterator is currently located at the last version for a key.
 func (i *MVCCIncrementalIterator) NextKey() {
 	i.iter.NextKey()
+	i.iteratorAdvanceCount++
 	if !i.checkValidAndSaveErr() {
 		return
 	}
@@ -287,6 +304,7 @@ func (i *MVCCIncrementalIterator) maybeSkipKeys() {
 			// expensive than a Next call.
 			seekKey := MakeMVCCMetadataKey(tbiKey)
 			i.iter.SeekGE(seekKey)
+			i.iteratorAdvanceCount++
 			if !i.checkValidAndSaveErr() {
 				return
 			}
@@ -367,11 +385,20 @@ func (i *MVCCIncrementalIterator) advance() {
 			return
 		}
 
+		// If we reached requested iteration limit here, stop iteration with an error and
+		// provide a resume key.
+		if i.maxAllowedIterations > 0 && i.iteratorAdvanceCount > i.maxAllowedIterations {
+			i.err = &ResourceLimitError{Limit: i.maxAllowedIterations, ResumeKey: i.Key()}
+			i.valid = false
+			return
+		}
+
 		// We have encountered an intent but it does not lie in the timestamp span
 		// (startTime, endTime] so we do not throw an error, and attempt to move to
 		// the next valid KV.
 		if i.meta.Txn != nil {
 			i.iter.Next()
+			i.iteratorAdvanceCount++
 			if !i.checkValidAndSaveErr() {
 				return
 			}
@@ -391,10 +418,8 @@ func (i *MVCCIncrementalIterator) advance() {
 			// done.
 			break
 		}
-
-		if ok, err := i.iter.Valid(); !ok {
-			i.err = err
-			i.valid = false
+		i.iteratorAdvanceCount++
+		if !i.checkValidAndSaveErr() {
 			return
 		}
 	}
@@ -439,9 +464,8 @@ func (i *MVCCIncrementalIterator) UnsafeValue() []byte {
 func (i *MVCCIncrementalIterator) NextIgnoringTime() {
 	for {
 		i.iter.Next()
-		if ok, err := i.iter.Valid(); !ok {
-			i.err = err
-			i.valid = false
+		i.iteratorAdvanceCount++
+		if !i.checkValidAndSaveErr() {
 			return
 		}
 
@@ -479,3 +503,18 @@ func (i *MVCCIncrementalIterator) TryGetIntentError() error {
 		Intents: i.intents,
 	}
 }
+
+// ResourceLimitError is returned by MVCCIncrementalIterator when iterator reaches
+// maximum number of iterations in the underlying storage iterator. Error will
+// provide used Limit as well as ResumeKey that could be used to resume iteration
+// on the subsequent attempt.
+type ResourceLimitError struct {
+	Limit     int64
+	ResumeKey MVCCKey
+}
+
+var _ error = &ResourceLimitError{}
+
+func (e *ResourceLimitError) Error() string {
+	return fmt.Sprintf("iteration count exhausted limit %d", e.Limit)
+}