batcheval: support MaxSpanRequestKeys and TargetBytes for FlushLockTable

This allows the caller to limit how many keys might be flushed from the
lock table.

Informs #146356
Release note: None

Author: stevendanna

pkg/kv/kvclient/kvcoord/dist_sender.go

@@ -1101,7 +1101,9 @@ func (ds *DistSender) initAndVerifyBatch(ctx context.Context, ba *kvpb.BatchRequ
 			*kvpb.GetRequest, *kvpb.ResolveIntentRequest, *kvpb.DeleteRequest, *kvpb.PutRequest:
 			// Accepted point requests that can be in batches with limit. No
 			// need to set disallowedReq.
-
+		case *kvpb.FlushLockTableRequest:
+			// FlushLockTableRequest handles limits itself. It is also an isAlone request so
+			// most of these checks should not be relevant.
 		default:
 			disallowedReq = inner.Method().String()
 		}

pkg/kv/kvpb/api.proto

@@ -2802,6 +2802,7 @@ message Header {
   // - ResolveIntentRangeRequest
   // - QueryLocksRequest
   // - IsSpanEmptyRequest
+  // - FlushLockTableRequest
   //
   // The following request types are also allowed in the batch. These requests
   // do not consume any keys from the limit; however, they may fail to be

pkg/kv/kvserver/batcheval/cmd_flush_lock_table.go

@@ -44,13 +44,34 @@ func declareKeysFlushLockTable(
 func FlushLockTable(
 	ctx context.Context, rw storage.ReadWriter, cArgs CommandArgs, response kvpb.Response,
 ) (result.Result, error) {
+	header := cArgs.Header
 	args := cArgs.Args.(*kvpb.FlushLockTableRequest)
 	resp := response.(*kvpb.FlushLockTableResponse)
 
-	// TODO(ssd): Allow the caller to limit how many locks we write out.
-	locksToFlush := make([]roachpb.LockAcquisition, 0)
-	cArgs.EvalCtx.GetConcurrencyManager().ExportUnreplicatedLocks(args.Span(), func(l *roachpb.LockAcquisition) {
+	var (
+		resumeSpan   *roachpb.Span
+		resumeReason kvpb.ResumeReason
+		currentSize  int64
+
+		locksToFlush = make([]roachpb.LockAcquisition, 0)
+	)
+	cArgs.EvalCtx.GetConcurrencyManager().ExportUnreplicatedLocks(args.Span(), func(l *roachpb.LockAcquisition) bool {
+		if max := header.MaxSpanRequestKeys; max > 0 && int64(len(locksToFlush)) >= max {
+			resumeReason = kvpb.RESUME_KEY_LIMIT
+			resumeSpan = &roachpb.Span{Key: args.Key, EndKey: l.Key}
+			return false
+		}
+		if maxSize := header.TargetBytes; maxSize > 0 {
+			sz := storage.ApproximateLockTableSize(l)
+			if currentSize+sz > maxSize {
+				resumeReason = kvpb.RESUME_KEY_LIMIT
+				resumeSpan = &roachpb.Span{Key: args.Key, EndKey: l.Key}
+				return false
+			}
+			currentSize += sz
+		}
 		locksToFlush = append(locksToFlush, *l)
+		return true
 	})
 
 	for i, l := range locksToFlush {
@@ -61,6 +82,9 @@ func FlushLockTable(
 			return result.Result{}, err
 		}
 	}
+
+	resp.ResumeReason = resumeReason
+	resp.ResumeSpan = resumeSpan
 	resp.LocksWritten = uint64(len(locksToFlush))
 
 	// NOTE: The locks still exist in the in-memory lock table. They are not

pkg/kv/kvserver/client_lock_table_test.go

@@ -163,6 +163,8 @@ func TestClientLockTableDataDriven(t *testing.T) {
 				startKey := evalCtx.getNamedKey("start", d)
 				endKey := evalCtx.getNamedKey("end", d)
 				b := db.NewBatch()
+				b.Header.MaxSpanRequestKeys = int64(evalCtx.getInt(d, "max-keys"))
+				b.Header.TargetBytes = int64(evalCtx.getInt(d, "target-bytes"))
 				b.AddRawRequest(&kvpb.FlushLockTableRequest{
 					RequestHeader: kvpb.RequestHeader{
 						Key:    startKey,
@@ -304,6 +306,15 @@ func (e *evalCtx) getLockWaitPolicy(d *datadriven.TestData) lock.WaitPolicy {
 	}
 }
 
+func (e *evalCtx) getInt(d *datadriven.TestData, name string) int {
+	if d.HasArg(name) {
+		var i int
+		d.ScanArgs(e.t, name, &i)
+		return i
+	}
+	return 0
+}
+
 func (e *evalCtx) cmdInBatch(cmdStr string, b *kv.Batch) error {
 	d := datadriven.TestData{}
 	var err error

pkg/kv/kvserver/concurrency/concurrency_control.go

@@ -269,8 +269,8 @@ type LockManager interface {
 	QueryLockTableState(ctx context.Context, span roachpb.Span, opts QueryLockTableOptions) ([]roachpb.LockStateInfo, QueryLockTableResumeState)
 
 	// ExportUnreplicatedLocks runs exporter on each held, unreplicated lock
-	// in the given span.
-	ExportUnreplicatedLocks(span roachpb.Span, exporter func(*roachpb.LockAcquisition))
+	// in the given span until the exporter returns false.
+	ExportUnreplicatedLocks(span roachpb.Span, exporter func(*roachpb.LockAcquisition) bool)
 }
 
 // TransactionManager is concerned with tracking transactions that have their
@@ -776,12 +776,12 @@ type lockTable interface {
 	QueryLockTableState(span roachpb.Span, opts QueryLockTableOptions) ([]roachpb.LockStateInfo, QueryLockTableResumeState)
 
 	// ExportUnreplicatedLocks runs exporter on each held, unreplicated lock
-	// in the given span.
+	// in the given span until the exporter returns false.
 	//
 	// Note that the caller is responsible for acquiring latches across the span
 	// it is exporting if it needs to be sure that the exported locks won't be
 	// updated in the lock table while it is still referencing them.
-	ExportUnreplicatedLocks(span roachpb.Span, exporter func(*roachpb.LockAcquisition))
+	ExportUnreplicatedLocks(span roachpb.Span, exporter func(*roachpb.LockAcquisition) bool)
 
 	// Metrics returns information about the state of the lockTable.
 	Metrics() LockTableMetrics

pkg/kv/kvserver/concurrency/concurrency_manager.go

@@ -619,7 +619,9 @@ func (m *managerImpl) QueryLockTableState(
 }
 
 // ExportUnreplicatedLocks implements the LockManager interface.
-func (m *managerImpl) ExportUnreplicatedLocks(span roachpb.Span, f func(*roachpb.LockAcquisition)) {
+func (m *managerImpl) ExportUnreplicatedLocks(
+	span roachpb.Span, f func(*roachpb.LockAcquisition) bool,
+) {
 	m.lt.ExportUnreplicatedLocks(span, f)
 }
 
@@ -734,9 +736,10 @@ func (m *managerImpl) exportUnreplicatedLocks() ([]*roachpb.LockAcquisition, int
 	// TODO(ssd): Expose a function that allows us to pre-allocate this a bit better.
 	approximateBatchSize := int64(0)
 	acquistions := make([]*roachpb.LockAcquisition, 0)
-	m.lt.ExportUnreplicatedLocks(allKeysSpan, func(acq *roachpb.LockAcquisition) {
+	m.lt.ExportUnreplicatedLocks(allKeysSpan, func(acq *roachpb.LockAcquisition) bool {
 		approximateBatchSize += storage.ApproximateLockTableSize(acq)
 		acquistions = append(acquistions, acq)
+		return true
 	})
 	return acquistions, approximateBatchSize
 }

pkg/kv/kvserver/concurrency/lock_table.go

@@ -4746,7 +4746,7 @@ func (t *lockTableImpl) ClearGE(key roachpb.Key) []roachpb.LockAcquisition {
 // TODO(ssd): Once we have the full set of functions we need for lock flushing, we
 // should do a refactoring pass to reduce some of the duplication.
 func (t *lockTableImpl) ExportUnreplicatedLocks(
-	span roachpb.Span, exporter func(*roachpb.LockAcquisition),
+	span roachpb.Span, exporter func(*roachpb.LockAcquisition) bool,
 ) {
 	t.enabledMu.RLock()
 	defer t.enabledMu.RUnlock()
@@ -4757,12 +4757,12 @@ func (t *lockTableImpl) ExportUnreplicatedLocks(
 	t.locks.mu.RLock()
 	defer t.locks.mu.RUnlock()
 
-	exportKeyLocks := func(l *keyLocks) {
+	exportKeyLocks := func(l *keyLocks) bool {
 		l.mu.Lock()
 		defer l.mu.Unlock()
 
 		if !l.key.Less(span.EndKey) {
-			return
+			return false
 		}
 
 		for hl := l.holders.Front(); hl != nil; hl = hl.Next() {
@@ -4777,7 +4777,7 @@ func (t *lockTableImpl) ExportUnreplicatedLocks(
 
 			for _, str := range unreplicatedHolderStrengths {
 				if tl.unreplicatedInfo.held(str) {
-					exporter(&roachpb.LockAcquisition{
+					keepGoing := exporter(&roachpb.LockAcquisition{
 						Span: roachpb.Span{
 							Key: l.key,
 						},
@@ -4786,14 +4786,20 @@ func (t *lockTableImpl) ExportUnreplicatedLocks(
 						Strength:       str,
 						IgnoredSeqNums: tl.unreplicatedInfo.ignoredSeqNums,
 					})
+					if !keepGoing {
+						return false
+					}
 				}
 			}
 		}
+		return true
 	}
 
 	iter := t.locks.MakeIter()
 	for iter.SeekGE(&keyLocks{key: span.Key}); iter.Valid(); iter.Next() {
-		exportKeyLocks(iter.Cur())
+		if !exportKeyLocks(iter.Cur()) {
+			break
+		}
 	}
 }
 

pkg/kv/kvserver/concurrency/verifiable_lock_table.go

@@ -118,7 +118,7 @@ func (v verifyingLockTable) QueryLockTableState(
 }
 
 func (v verifyingLockTable) ExportUnreplicatedLocks(
-	span roachpb.Span, exporter func(*roachpb.LockAcquisition),
+	span roachpb.Span, exporter func(*roachpb.LockAcquisition) bool,
 ) {
 	v.lt.ExportUnreplicatedLocks(span, exporter)
 }

pkg/kv/kvserver/testdata/lock_table/flush

@@ -5,6 +5,9 @@ new-txn txn=t1
 put txn=t1 k=a v=bar
 ----
 
+put txn=t1 k=b v=baz
+----
+
 commit txn=t1
 ----
 
@@ -34,3 +37,145 @@ key: "\xfaa", str: Exclusive, txn: t2
 
 commit txn=t2
 ----
+
+# MaxSpanRequestKeys Limit
+new-txn txn=t3
+----
+
+get txn=t3 k=a lock=Exclusive dur=Unreplicated
+----
+get: "\xfaa"="bar"
+
+get txn=t3 k=b lock=Exclusive dur=Unreplicated
+----
+get: "\xfab"="baz"
+
+print-in-memory-lock-table
+----
+num=2
+ lock: "\xfaa"
+  holder: txn: t3 epoch: 0, iso: Serializable, ts: <stripped>, info: unrepl [(str: Exclusive seq: 0)]
+ lock: "\xfab"
+  holder: txn: t3 epoch: 0, iso: Serializable, ts: <stripped>, info: unrepl [(str: Exclusive seq: 0)]
+
+flush-lock-table start=a end=z max-keys=1
+----
+
+print-in-memory-lock-table
+----
+num=1
+ lock: "\xfab"
+  holder: txn: t3 epoch: 0, iso: Serializable, ts: <stripped>, info: unrepl [(str: Exclusive seq: 0)]
+
+print-replicated-lock-table start=a end=z
+----
+key: "\xfaa", str: Exclusive, txn: t3
+
+commit txn=t3
+----
+
+# TargetBytes Limit
+new-txn txn=t4
+----
+
+get txn=t4 k=a lock=Exclusive dur=Unreplicated
+----
+get: "\xfaa"="bar"
+
+get txn=t4 k=b lock=Exclusive dur=Unreplicated
+----
+get: "\xfab"="baz"
+
+print-in-memory-lock-table
+----
+num=2
+ lock: "\xfaa"
+  holder: txn: t4 epoch: 0, iso: Serializable, ts: <stripped>, info: unrepl [(str: Exclusive seq: 0)]
+ lock: "\xfab"
+  holder: txn: t4 epoch: 0, iso: Serializable, ts: <stripped>, info: unrepl [(str: Exclusive seq: 0)]
+
+flush-lock-table start=a end=z target-bytes=128
+----
+
+print-in-memory-lock-table
+----
+num=1
+ lock: "\xfab"
+  holder: txn: t4 epoch: 0, iso: Serializable, ts: <stripped>, info: unrepl [(str: Exclusive seq: 0)]
+
+print-replicated-lock-table start=a end=z
+----
+key: "\xfaa", str: Exclusive, txn: t4
+
+commit txn=t4
+----
+
+# MaxSpanRequestKeys ineffective limit
+new-txn txn=t5
+----
+
+get txn=t5 k=a lock=Exclusive dur=Unreplicated
+----
+get: "\xfaa"="bar"
+
+get txn=t5 k=b lock=Exclusive dur=Unreplicated
+----
+get: "\xfab"="baz"
+
+print-in-memory-lock-table
+----
+num=2
+ lock: "\xfaa"
+  holder: txn: t5 epoch: 0, iso: Serializable, ts: <stripped>, info: unrepl [(str: Exclusive seq: 0)]
+ lock: "\xfab"
+  holder: txn: t5 epoch: 0, iso: Serializable, ts: <stripped>, info: unrepl [(str: Exclusive seq: 0)]
+
+flush-lock-table start=a end=z max-keys=8
+----
+
+print-in-memory-lock-table
+----
+num=0
+
+print-replicated-lock-table start=a end=z
+----
+key: "\xfaa", str: Exclusive, txn: t5
+key: "\xfab", str: Exclusive, txn: t5
+
+commit txn=t5
+----
+
+# TargetBytes ineffective limit
+new-txn txn=t6
+----
+
+get txn=t6 k=a lock=Exclusive dur=Unreplicated
+----
+get: "\xfaa"="bar"
+
+get txn=t6 k=b lock=Exclusive dur=Unreplicated
+----
+get: "\xfab"="baz"
+
+print-in-memory-lock-table
+----
+num=2
+ lock: "\xfaa"
+  holder: txn: t6 epoch: 0, iso: Serializable, ts: <stripped>, info: unrepl [(str: Exclusive seq: 0)]
+ lock: "\xfab"
+  holder: txn: t6 epoch: 0, iso: Serializable, ts: <stripped>, info: unrepl [(str: Exclusive seq: 0)]
+
+flush-lock-table start=a end=z target-bytes=4096
+----
+
+print-in-memory-lock-table
+----
+num=0
+
+print-replicated-lock-table start=a end=z
+----
+key: "\xfaa", str: Exclusive, txn: t6
+key: "\xfab", str: Exclusive, txn: t6
+
+commit txn=t6
+----