restore: reduce allocations when restoring to longer tenant keys

Release note: none. Epic: none.
cockroachdb · Mar 8, 2024 · 2bef9e2 · 2bef9e2
1 parent a4f061d
commit 2bef9e2
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/pkg/ccl/backupccl/key_rewriter.go b/pkg/ccl/backupccl/key_rewriter.go
@@ -94,6 +94,7 @@ type KeyRewriter struct {
 	prefixes prefixRewriter
 	tenants  prefixRewriter
 	descs    map[descpb.ID]catalog.TableDescriptor
+	alloc    []byte
 }
 
 // MakeKeyRewriterFromRekeys makes a KeyRewriter from Rekey protos.
@@ -273,7 +274,19 @@ func (kr *KeyRewriter) RewriteKey(
 		copy(keyTenantPrefix, newTenantPrefix)
 		rekeyed = append(keyTenantPrefix, rekeyed...)
 	} else {
-		rekeyed = append(newTenantPrefix, rekeyed...)
+		l := len(newTenantPrefix) + len(rekeyed)
+		if len(kr.alloc) < l {
+			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