backfill: split merge memory accounting mutex acquisitions

Previously, constructMergeBatch() would get the mutex on memory
accounting once per batch to minimize the number of atomics needed to
perform memory accounting. With the addition of vector indexing, we now
need to perform KV operations in order to construct batches, which can
lead to deadlocks between KV and the memory accounting mutex.

This patch splits the mutex acquisition into two acquisitions, releasing
the mutex in between. Memory used is accumulated as merged entries are
built and then accounting is notified for the entire batch. This has the
effect of making the accounting lag behind allocation a bit, but that
seems preferable than paying the cost of a mutex acquisition for every
row merged.

Fixes: #155190
Release note (bug fix): A potential deadlock during vector index
creation has been corrected.

Author: mw5h

pkg/sql/backfill/mvcc_index_merger.go

@@ -518,28 +518,32 @@ func (ibm *IndexBackfillMerger) constructMergeBatch(
 		return nil, 0, 0, 0, err
 	}
 
-	// We acquire the bound account lock for the entirety of the merge batch
-	// construction so that we don't have to acquire the lock every time we want
-	// to grow the account.
 	var memUsedInMerge int64
-	ibm.muBoundAccount.Lock()
-	defer ibm.muBoundAccount.Unlock()
-	for i := range rb.Results {
-		// Since the source index is delete-preserving, reading the latest value for
-		// a key that has existed in the past should always return a value.
-		if rb.Results[i].Rows[0].Value == nil {
-			return nil, 0, 0, 0, errors.AssertionFailedf("expected value to be present in temp index for key=%s", rb.Results[i].Rows[0].Key)
-		}
-		rowMem := int64(len(rb.Results[i].Rows[0].Key)) + int64(len(rb.Results[i].Rows[0].Value.RawBytes))
-		if err := ibm.muBoundAccount.boundAccount.Grow(ctx, rowMem); err != nil {
-			return nil, 0, 0, 0, errors.Wrap(err, "failed to allocate space to read latest keys from temp index")
+	err := func() error {
+		ibm.muBoundAccount.Lock()
+		defer ibm.muBoundAccount.Unlock()
+		for i := range rb.Results {
+			// Since the source index is delete-preserving, reading the latest value for
+			// a key that has existed in the past should always return a value.
+			if rb.Results[i].Rows[0].Value == nil {
+				return errors.AssertionFailedf("expected value to be present in temp index for key=%s", rb.Results[i].Rows[0].Key)
+			}
+			rowMem := int64(len(rb.Results[i].Rows[0].Key)) + int64(len(rb.Results[i].Rows[0].Value.RawBytes))
+			if err := ibm.muBoundAccount.boundAccount.Grow(ctx, rowMem); err != nil {
+				return errors.Wrap(err, "failed to allocate space to read latest keys from temp index")
+			}
+			memUsedInMerge += rowMem
 		}
-		memUsedInMerge += rowMem
+		return nil
+	}()
+	if err != nil {
+		return nil, 0, 0, 0, err
 	}
 
 	prefixLen := len(sourcePrefix)
 	destKey := make([]byte, len(destPrefix))
 	var deletedCount int
+	var totalEntryBytes int64
 	wb := txn.NewBatch()
 	resultOffset := 0
 	for sourceOffset := range sourceKeys {
@@ -571,18 +575,26 @@ func (ibm *IndexBackfillMerger) constructMergeBatch(
 			continue
 		}
 
-		entryBytes := mergedEntryBytes(mergedEntry, deleted)
-		if err := ibm.muBoundAccount.boundAccount.Grow(ctx, entryBytes); err != nil {
-			return nil, 0, 0, 0, errors.Wrap(err, "failed to allocate space to merge entry from temp index")
-		}
-		memUsedInMerge += entryBytes
+		totalEntryBytes += mergedEntryBytes(mergedEntry, deleted)
 		if deleted {
 			deletedCount++
 			wb.Del(mergedEntry.Key)
 		} else {
 			wb.Put(mergedEntry.Key, mergedEntry.Value)
 		}
 	}
+	err = func() error {
+		ibm.muBoundAccount.Lock()
+		defer ibm.muBoundAccount.Unlock()
+		if err := ibm.muBoundAccount.boundAccount.Grow(ctx, totalEntryBytes); err != nil {
+			return errors.Wrap(err, "failed to allocate space to merge entry from temp index")
+		}
+		return nil
+	}()
+	if err != nil {
+		return nil, 0, 0, 0, err
+	}
+	memUsedInMerge += totalEntryBytes
 
 	return wb, memUsedInMerge, deletedCount, keysToSkip, nil
 }