From e7316d642c46ebe3d4373149e16ef8c610bc7e3c Mon Sep 17 00:00:00 2001
From: David Taylor <tinystatemachine@gmail.com>
Date: Fri, 8 Mar 2024 19:16:45 +0000
Subject: [PATCH] restore: reduce allocations when restoring to longer tenant
 keys

Release note: none.
Epic: none.
---
 pkg/ccl/backupccl/key_rewriter.go | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/pkg/ccl/backupccl/key_rewriter.go b/pkg/ccl/backupccl/key_rewriter.go
index f126549a30c6..56165fad8291 100644
--- a/pkg/ccl/backupccl/key_rewriter.go
+++ b/pkg/ccl/backupccl/key_rewriter.go
@@ -94,6 +94,9 @@ type KeyRewriter struct {
 	prefixes prefixRewriter
 	tenants  prefixRewriter
 	descs    map[descpb.ID]catalog.TableDescriptor
+	// alloc is used to amortize the cost of many small allocations for keys which
+	// change length during rewriting, preventing doing so in-place.
+	alloc []byte
 }
 
 // MakeKeyRewriterFromRekeys makes a KeyRewriter from Rekey protos.
@@ -273,7 +276,25 @@ func (kr *KeyRewriter) RewriteKey(
 		copy(keyTenantPrefix, newTenantPrefix)
 		rekeyed = append(keyTenantPrefix, rekeyed...)
 	} else {
-		rekeyed = append(newTenantPrefix, rekeyed...)
+		// Prefix length changed so we cannot rewrite it in-place; instead allocate
+		// the new key off the allocation slab, refilling it if needed first.
+		l := len(newTenantPrefix) + len(rekeyed)
+		if len(kr.alloc) < l {
+			// If individual keys are huge, allocate only 16 at a time, so that they
+			// can be freed sooner (aliasing) and one key doesn't massively over-alloc
+			// by a factor of 256. Otherwise allocate 256 worth so we don't need to
+			// do this again for awhile.
+			if l > 1<<20 {
+				kr.alloc = make([]byte, l*16)
+			} else {
+				kr.alloc = make([]byte, l*256)
+			}
+		}
+		tmp := kr.alloc[:l:l]
+		kr.alloc = kr.alloc[l:]
+		copy(tmp, newTenantPrefix)
+		copy(tmp[len(newTenantPrefix):], rekeyed)
+		rekeyed = tmp
 	}
 
 	return rekeyed, ok, err