Merge #148974

148974: util/bufalloc: fix misuse and improve ByteAllocator API r=yuzefovich a=yuzefovich

**kvserver/gc,rowcontainer: fix misuse of ByteAllocator API**

`bufalloc.ByteAllocator.Copy` has an option to allocate extra capacity,
meaning that it copies the passed-in byte slice into a new memory region
while also adding extra capacity to be utilized by caller. Crucially,
that extra capacity lives in the returned byte slice containing the
copy. We use this feature in two places throughout the codebase, and in
both places it was used incorrectly, resulting in a wasteful allocation.

Consider the following setup:

  k = [1, 2], v = [3, 4, 5, 6]

With invocation like `alloc, k = alloc.Copy(k, len(v))`, we would get

  len(k) = 2, cap(k) = 6

However, in both spots we copied `v` via `alloc, v = alloc.Copy(v, 0)`,
which would actually create another allocation.

This commit fixes this oversight (both were introduced in 2022 by different
authors), and the following commit will clean up the API to avoid possible
confusion.

Informs: #147601.

**util/bufalloc: clean up ByteAllocator API**

`bufalloc.ByteAllocator`, as evidenced by a fix in the previous commit,
has confusing API for allocating "extra capacity". This commit cleans it up
by removing the option altogether. The only two spots that utilized the
extra capacity previously now explicitly request combined allocated
upfront and manage the copying behavior.

Release note: None

Co-authored-by: Yahor Yuzefovich <[email protected]>

Author: craig[bot]

pkg/BUILD.bazel

@@ -713,6 +713,7 @@ ALL_TESTS = [
     "//pkg/util/binfetcher:binfetcher_test",
     "//pkg/util/bitarray:bitarray_test",
     "//pkg/util/bitmap:bitmap_test",
+    "//pkg/util/bufalloc:bufalloc_test",
     "//pkg/util/buildutil:buildutil_test",
     "//pkg/util/cache:cache_test",
     "//pkg/util/caller:caller_test",
@@ -2534,6 +2535,7 @@ GO_TARGETS = [
     "//pkg/util/bitmap:bitmap",
     "//pkg/util/bitmap:bitmap_test",
     "//pkg/util/bufalloc:bufalloc",
+    "//pkg/util/bufalloc:bufalloc_test",
     "//pkg/util/buildutil:buildutil",
     "//pkg/util/buildutil:buildutil_test",
     "//pkg/util/bulk:bulk",

pkg/ccl/changefeedccl/event_processing.go

@@ -453,14 +453,14 @@ func (c *kvEventToRowConsumer) encodeAndEmit(
 	if err != nil {
 		return err
 	}
-	c.scratch, keyCopy = c.scratch.Copy(encodedKey, 0 /* extraCap */)
+	c.scratch, keyCopy = c.scratch.Copy(encodedKey)
 	// TODO(yevgeniy): Some refactoring is needed in the encoder: namely, prevRow
 	// might not be available at all when working with changefeed expressions.
 	encodedValue, err := c.encoder.EncodeValue(ctx, evCtx, updatedRow, prevRow)
 	if err != nil {
 		return err
 	}
-	c.scratch, valueCopy = c.scratch.Copy(encodedValue, 0 /* extraCap */)
+	c.scratch, valueCopy = c.scratch.Copy(encodedValue)
 
 	// Since we're done processing/converting this event, and will not use much more
 	// than len(key)+len(bytes) worth of resources, adjust allocation to match.

pkg/ccl/changefeedccl/sink.go

@@ -494,7 +494,7 @@ func (s *bufferSink) EmitResolvedTimestamp(
 	if err != nil {
 		return err
 	}
-	s.scratch, payload = s.scratch.Copy(payload, 0 /* extraCap */)
+	s.scratch, payload = s.scratch.Copy(payload)
 	s.buf.Push(rowenc.EncDatumRow{
 		{Datum: tree.DNull}, // resolved span
 		{Datum: tree.DNull}, // topic

pkg/ccl/changefeedccl/sink_kafka.go

@@ -424,7 +424,7 @@ func (s *kafkaSink) EmitResolvedTimestamp(
 		if err != nil {
 			return err
 		}
-		s.scratch, payload = s.scratch.Copy(payload, 0 /* extraCap */)
+		s.scratch, payload = s.scratch.Copy(payload)
 
 		// sarama caches this, which is why we have to periodically refresh the
 		// metadata above. Staleness here does not impact correctness. Some new

pkg/ccl/changefeedccl/sink_sql.go

@@ -168,7 +168,7 @@ func (s *sqlSink) EmitResolvedTimestamp(
 		if err != nil {
 			return err
 		}
-		s.scratch, payload = s.scratch.Copy(payload, 0 /* extraCap */)
+		s.scratch, payload = s.scratch.Copy(payload)
 		for partition := int32(0); partition < sqlSinkNumPartitions; partition++ {
 			if err := s.emit(ctx, topic, partition, noKey, noValue, payload); err != nil {
 				return err

pkg/crosscluster/streamclient/randclient/random_stream_client.go

@@ -738,7 +738,7 @@ func duplicateEvent(event crosscluster.Event) crosscluster.Event {
 		for i := range kvs {
 			res[i].Key = kvs[i].KeyValue.Key.Clone()
 			res[i].Value.Timestamp = kvs[i].KeyValue.Value.Timestamp
-			a, res[i].Value.RawBytes = a.Copy(kvs[i].KeyValue.Value.RawBytes, 0)
+			a, res[i].Value.RawBytes = a.Copy(kvs[i].KeyValue.Value.RawBytes)
 		}
 		dup = crosscluster.MakeKVEventFromKVs(res)
 	case crosscluster.SSTableEvent:

pkg/kv/kvnemesis/engine.go

@@ -93,7 +93,7 @@ func (e *Engine) Get(key roachpb.Key, ts hlc.Timestamp) roachpb.Value {
 	if err != nil {
 		panic(err)
 	}
-	e.b, valCopy = e.b.Copy(v, 0 /* extraCap */)
+	e.b, valCopy = e.b.Copy(v)
 	mvccVal, err := storage.DecodeMVCCValue(valCopy)
 	if err != nil {
 		panic(err)
@@ -137,12 +137,12 @@ func (e *Engine) Iterate(
 	for iter.First(); iter.Valid(); iter.Next() {
 		hasPoint, _ := iter.HasPointAndRange()
 		var keyCopy, valCopy []byte
-		e.b, keyCopy = e.b.Copy(iter.Key(), 0 /* extraCap */)
+		e.b, keyCopy = e.b.Copy(iter.Key())
 		v, err := iter.ValueAndErr()
 		if err != nil {
 			fn(nil, nil, hlc.Timestamp{}, nil, err)
 		}
-		e.b, valCopy = e.b.Copy(v, 0 /* extraCap */)
+		e.b, valCopy = e.b.Copy(v)
 		if hasPoint {
 			key, err := storage.DecodeMVCCKey(keyCopy)
 			if err != nil {
@@ -153,8 +153,8 @@ func (e *Engine) Iterate(
 		}
 		if iter.RangeKeyChanged() {
 			keyCopy, endKeyCopy := iter.RangeBounds()
-			e.b, keyCopy = e.b.Copy(keyCopy, 0 /* extraCap */)
-			e.b, endKeyCopy = e.b.Copy(endKeyCopy, 0 /* extraCap */)
+			e.b, keyCopy = e.b.Copy(keyCopy)
+			e.b, endKeyCopy = e.b.Copy(endKeyCopy)
 			for _, rk := range iter.RangeKeys() {
 				ts, err := mvccencoding.DecodeMVCCTimestampSuffix(rk.Suffix)
 				if err != nil {
@@ -170,7 +170,7 @@ func (e *Engine) Iterate(
 					fn(nil, nil, hlc.Timestamp{}, nil, errors.Errorf("invalid key %q", endKeyCopy))
 				}
 
-				e.b, rk.Value = e.b.Copy(rk.Value, 0)
+				e.b, rk.Value = e.b.Copy(rk.Value)
 
 				fn(engineKey.Key, engineEndKey.Key, ts, rk.Value, nil)
 			}

pkg/kv/kvserver/gc/gc.go

@@ -809,7 +809,7 @@ func (b *gcKeyBatcher) foundGarbage(
 		// Whenever new key is started or new batch is started with the same key in
 		// it, record key value using batches' allocator.
 		if b.prevWasNewest || len(b.pointsBatches[i].batchGCKeys) == 0 {
-			b.pointsBatches[i].alloc, key = b.pointsBatches[i].alloc.Copy(cur.key.Key, 0)
+			b.pointsBatches[i].alloc, key = b.pointsBatches[i].alloc.Copy(cur.key.Key)
 			b.pointsBatches[i].batchGCKeys = append(b.pointsBatches[i].batchGCKeys,
 				kvpb.GCRequest_GCKey{Key: key, Timestamp: cur.key.Timestamp})
 			keyMemUsed := len(key) + hlcTimestampSize
@@ -1065,7 +1065,7 @@ func (b *intentBatcher) addAndMaybeFlushIntents(
 	// We need to register passed intent regardless of flushing operation result
 	// so that batcher is left in consistent state and don't miss any keys if
 	// caller resumes batching.
-	b.alloc, key = b.alloc.Copy(key, 0)
+	b.alloc, key = b.alloc.Copy(key)
 	b.pendingLocks = append(b.pendingLocks, roachpb.MakeLock(meta.Txn, key, str))
 	b.collectedIntentBytes += int64(len(key))
 	b.pendingTxns[txnID] = true

pkg/kv/kvserver/gc/gc_iterator.go

@@ -237,10 +237,12 @@ func (b *gcIteratorRingBuf) pushBack(
 	}
 	i := (b.head + b.len) % gcIteratorRingBufSize
 	b.allocs[i] = b.allocs[i].Truncate()
-	b.allocs[i], k.Key = b.allocs[i].Copy(k.Key, len(metaValue))
-	if len(metaValue) > 0 {
-		b.allocs[i], metaValue = b.allocs[i].Copy(metaValue, 0)
-	}
+	var buf []byte
+	b.allocs[i], buf = b.allocs[i].Alloc(len(k.Key) + len(metaValue))
+	copy(buf, k.Key)
+	k.Key, buf = buf[:len(k.Key):len(k.Key)], buf[len(k.Key):]
+	copy(buf, metaValue)
+	metaValue = buf
 	b.buf[i] = mvccKeyValue{
 		key:                  k,
 		mvccValueLen:         mvccValueLen,

pkg/kv/kvserver/rangefeed/catchup_scan.go

@@ -198,8 +198,8 @@ func (i *CatchUpIterator) CatchUpScan(
 				rangeKeys := i.RangeKeys()
 				for j := rangeKeys.Len() - 1; j >= 0; j-- {
 					var span roachpb.Span
-					a, span.Key = a.Copy(rangeKeys.Bounds.Key, 0)
-					a, span.EndKey = a.Copy(rangeKeys.Bounds.EndKey, 0)
+					a, span.Key = a.Copy(rangeKeys.Bounds.Key)
+					a, span.EndKey = a.Copy(rangeKeys.Bounds.EndKey)
 					ts := rangeKeys.Versions[j].Timestamp
 					err := outputFn(&kvpb.RangeFeedEvent{
 						DeleteRange: &kvpb.RangeFeedDeleteRange{
@@ -298,7 +298,7 @@ func (i *CatchUpIterator) CatchUpScan(
 			if err := outputEvents(); err != nil {
 				return err
 			}
-			a, lastKey = a.Copy(unsafeKey.Key, 0)
+			a, lastKey = a.Copy(unsafeKey.Key)
 		}
 		key := lastKey
 
@@ -316,7 +316,7 @@ func (i *CatchUpIterator) CatchUpScan(
 		//   value.
 		if !ignore || (withDiff && len(reorderBuf) > 0) {
 			var val []byte
-			a, val = a.Copy(unsafeVal, 0)
+			a, val = a.Copy(unsafeVal)
 			if withDiff {
 				// Update the last version with its previous value (this version).
 				if l := len(reorderBuf) - 1; l >= 0 {