Merge #74735

74735: util/mon: add optional threadsafety to BoundAccount r=dt a=dt

BoundAccount is often passed around in the implementation of a processor
or task along side the maps, slices, etc that it tracks so that it can
be Grow()ed and Shrink()ed as they are mutated.

However, while most SQL processors and the libraries they use are all
single threaded, many bulk and cdc processors are not. The fact that
boundAccount is not threadsafe can then make it more difficult to use
in these areas. A wrapper that adds threadsafety could help but this
is difficult to use as functions that take a BoundAccount by name
cannot be passed a wrapper, since it is not an interface. Another
approach is to make a new boundAccount for each goroutine, but this
requires then careful merging of each of those back to the account
that was originally passed.

Instead, this change adds _optional_ threadsafety to mon.BoundAccount
by allowing caller to configure a field of it with a given mutex. When
this field is non-nil, the usage methods like Grow, Shrink, and Resize
will acquire and release it when called, making them threadsafe.
Lifecycle methods, like Init, Close or Empty are not included here.

Checking for this mutex in existing, single-threaded users will add some
additional overhead; BenchmarkBoundAccountGrow shows this to be about
1.2ns on my machine.

name                 old time/op  new time/op  delta
BoundAccountGrow-10  2.21ns ± 4%  3.42ns ± 3%  +55.23%  (p=0.000 n=10+9)

Release note: none.

Co-authored-by: David Taylor <[email protected]>

pkg/util/mon/bytes_usage.go

@@ -471,17 +471,25 @@ func (mm *BytesMonitor) Resource() Resource {
 // at once when it completes its work. Internally, BoundAccount amortizes
 // allocations from whichever BoundAccount it is associated with by allocating
 // additional memory and parceling it out (see BoundAccount.reserved). A nil
-// BoundAccount acts as an unlimited account for which growing and shrinking
-// are noops.
+// BoundAccount acts as an unlimited account for which growing and shrinking are
+// noops.
 //
 // See the comments in bytes_usage.go for a fuller picture of how these accounts
 // are used in CockroachDB.
+//
+// A normal BoundAccount is not safe for concurrent use by multiple goroutines,
+// however if the Mu field is set to a non-nil mutex, some methods such as Grow,
+// Shrink, and Resize calls will lock and unlock that mutex making them safe;
+// such methods are identified in their comments.
 type BoundAccount struct {
 	used int64
 	// reserved is a small buffer to amortize the cost of growing an account. It
 	// decreases as used increases (and vice-versa).
 	reserved int64
 	mon      *BytesMonitor
+
+	// Mu, if non-nil, is used in some methods such as Grow and Shrink.
+	Mu *syncutil.Mutex
 }
 
 // MakeStandaloneBudget creates a BoundAccount suitable for root
@@ -491,10 +499,15 @@ func MakeStandaloneBudget(capacity int64) BoundAccount {
 }
 
 // Used returns the number of bytes currently allocated through this account.
+// If Mu is set, it is safe for use by concurrent goroutines.
 func (b *BoundAccount) Used() int64 {
 	if b == nil {
 		return 0
 	}
+	if b.Mu != nil {
+		b.Mu.Lock()
+		defer b.Mu.Unlock()
+	}
 	return b.used
 }
 
@@ -585,10 +598,16 @@ func (b *BoundAccount) Close(ctx context.Context) {
 // If one is interested in specifying the new size of the account as a whole (as
 // opposed to resizing one object among many in the account), ResizeTo() should
 // be used.
+//
+// If Mu is set, it is safe for use by concurrent goroutines.
 func (b *BoundAccount) Resize(ctx context.Context, oldSz, newSz int64) error {
 	if b == nil {
 		return nil
 	}
+	if b.Mu != nil {
+		b.Mu.Lock()
+		defer b.Mu.Unlock()
+	}
 	delta := newSz - oldSz
 	switch {
 	case delta > 0:
@@ -600,10 +619,16 @@ func (b *BoundAccount) Resize(ctx context.Context, oldSz, newSz int64) error {
 }
 
 // ResizeTo resizes (grows or shrinks) the account to a specified size.
+//
+// If Mu is set, it is safe for use by concurrent goroutines.
 func (b *BoundAccount) ResizeTo(ctx context.Context, newSz int64) error {
 	if b == nil {
 		return nil
 	}
+	if b.Mu != nil {
+		b.Mu.Lock()
+		defer b.Mu.Unlock()
+	}
 	if newSz == b.used {
 		// Performance optimization to avoid an unnecessary dispatch.
 		return nil
@@ -612,10 +637,16 @@ func (b *BoundAccount) ResizeTo(ctx context.Context, newSz int64) error {
 }
 
 // Grow is an accessor for b.mon.GrowAccount.
+//
+// If Mu is set, it is safe for use by concurrent goroutines.
 func (b *BoundAccount) Grow(ctx context.Context, x int64) error {
 	if b == nil {
 		return nil
 	}
+	if b.Mu != nil {
+		b.Mu.Lock()
+		defer b.Mu.Unlock()
+	}
 	if b.reserved < x {
 		minExtra := b.mon.roundSize(x)
 		if err := b.mon.reserveBytes(ctx, minExtra); err != nil {
@@ -629,10 +660,16 @@ func (b *BoundAccount) Grow(ctx context.Context, x int64) error {
 }
 
 // Shrink releases part of the cumulated allocations by the specified size.
+//
+// If Mu is set, it is safe for use by concurrent goroutines.
 func (b *BoundAccount) Shrink(ctx context.Context, delta int64) {
 	if b == nil {
 		return
 	}
+	if b.Mu != nil {
+		b.Mu.Lock()
+		defer b.Mu.Unlock()
+	}
 	if b.used < delta {
 		logcrash.ReportOrPanic(ctx, &b.mon.settings.SV,
 			"%s: no bytes in account to release, current %d, free %d",