TiDB BR 的 Key Rewrite 机制：跨 Schema 版本恢复与 SST 导入路径的源码级解析

本文基于 TiDB 仓库中 BR 的设计讨论文档 Key Rewrite 讨论，梳理 Key Rewrite 的设计动机、数据结构与正反向改写算法，并结合当前 rewrite_rule.go 等源码实现，说明这一机制如何支撑 BR 恢复到 Table ID 不同的集群，以及在 TiKV Ingest 导入路径中的具体落点。读完本文，你将掌握 BR 恢复流程中 Rewrite Rule 的生成、校验与应用方式，以及“在每个副本 ingest 前就地改写 Key”这一方案背后的性能权衡。

为什么需要 Key Rewrite

BR 将备份的 SST 文件恢复到目标集群时，会遇到一个根本问题：备份源与恢复目标的 Schema Version 可能不同。例如备份是在某个旧版本集群上做的，而目标集群上对应表被删除过又重建过，Table ID 发生了变化；此时备份 SST 中的 Key（以旧 Table ID 编码的 t«tid»_ 前缀开头）直接导入目标集群会写入错误的表。Key Rewrite 正是为此设计的。

按设计文档，Key Rewrite 的目的有二：

1. 为 BR 提供修改 Table ID 的功能，以支持恢复到 Schema Version 不同的集群；
2. 为 Lightning 提供添加前缀的功能，省略 Lightning 与 Importer 之间重复的数据传输。

两条约束贯穿整个设计：

• 一个 BR 的 SST 文件可能包含多个 Table 的数据，因此必须支持多条 Rewrite Rule 同时生效；
• SST 可能来自非 TiDB 系统，因此 Importer 不应该对 Key 编码格式做任何假设（Key 不一定是 t«tid»_ 开头），规则必须基于“前缀字节串”而非结构化的整数 ID 来表达。

Rewrite Rule 的数据结构

设计文档给出的核心数据结构如下：

message RewriteRule {
	bytes old_prefix = 1;  // this can be empty for universal prefix insertion!
	bytes new_prefix = 2;  // these are _not_ just an integer!
}

message RestoreRequest {
	...
	repeated RewriteRule rewrite_rules = N;
	...
}

几个要点：

• old_prefix 可以为空，空前缀规则用于“给所有 Key 统一添加新前缀”的场景（即文档提到的 universal prefix insertion，这正是 Lightning 加前缀用法的来源）；
• 两个字段都是 bytes 而非整数，强调规则是通用前缀替换，与 TiDB 的 Key 编码格式解耦；
• 一个恢复请求携带一组规则，对应“一个 SST 内含多个表”的实际情况。

正向改写与反向还原算法

文档给出了两条对称的算法。正向改写一个 Key（找到第一条前缀匹配的规则，替换前缀后返回）：

fn rewrite_key(rules: &[RewriteRule], key: &[u8]) -> Cow<[u8]> {
    for rule in rules {
        if key.starts_with(rule.old_prefix) {
            return Cow::Owned(rule.new_prefix + key[rule.old_prefix.len()..])
        }
    }
    Cow::Borrowed(key)
}

反向还原一个 Key（用于需要把改写后的 Key 映射回原始 Key 的场景，例如校验和比对、Range 反查）：

fn undo_rewrite_key(rules: &[RewriteRule], key: &[u8]) -> Cow<[u8]> {
    for rule in rules {
        if key.starts_with(rule.new_prefix) {
            return Cow::Owned(rule.old_prefix + key[rule.new_prefix.len()..])
        }
    }
    Cow::Borrowed(key)
}

两者都是“线性扫描 + 首个前缀命中即返回”，未命中时 Key 原样透传（Cow::Borrowed，零拷贝）。这个简单性不是巧合：正因为规则是纯前缀替换，改写不会破坏 Key 的全序关系——同一前缀下的 Key 相对顺序保持不变，改写后的 Range 依然连续，这是后面能在导入流程中安全地 Pre-split、按 Range 并行导入的前提。

当前源码中的规则结构与生成方式

当前 BR 的实现位于 br/pkg/restore/utils/rewrite_rule.go。规则容器在 import_sstpb.RewriteRule（kvproto 定义，字段为 OldKeyPrefix / NewKeyPrefix / NewTimestamp 等）之上扩展为 RewriteRules 结构：

// RewriteRules contains rules for rewriting keys of tables.
type RewriteRules struct {
	Data        []*import_sstpb.RewriteRule
	OldKeyspace []byte
	NewKeyspace []byte
	// used to record checkpoint data
	NewTableID int64

	ShiftStartTs uint64
	StartTs      uint64
	RestoredTs   uint64
	// used to record backup files to pitr.
	TableIDRemapHint []TableIDRemap
}

相比 2019 年的设计文档，可以看到两处自然的演进：

• Keyspace 迁移：OldKeyspace / NewKeyspace 字段支持跨 Keyspace 恢复，即恢复时同时替换 Keyspace 前缀。import.go 中定义了 RewriteMode（RewriteModeLegacy 表示不应用规则、RewriteModeKeyspace 表示规则作用于 Keyspace）来区分这两种模式；
• 增量恢复（PITR）的时间过滤：每条 RewriteRule 上的 NewTimestamp / IgnoreBeforeTimestamp / IgnoreAfterTimestamp 让 TiKV 在 ingest 时按时间戳过滤掉窗口外的事务数据。SetTimeRangeFilter 按列族区分取法：default 列族取 min(ShiftStartTs, StartTs) 作为下界（因为大事务的 start ts 可能早于小事务，用 StartTs 过滤更保守安全），write 列族直接取 StartTs。

规则由备份侧表结构与恢复侧新表结构的 ID 映射生成。GetRewriteRules 接收新旧 TableInfo，通过 GetTableIDMap / GetIndexIDMap 计算映射，再产出两种粒度的规则：

• 粗粒度（getDetailRule = false）：整表一条规则，EncodeTablePrefix(oldTableID) -> EncodeTablePrefix(newTableID)；
• 细粒度（getDetailRule = true，常规全量恢复使用，见 client.go 中 GetRewriteRules(newTableInfo, table.Info, newTS, true) 的调用）：数据行一条规则（GenTableRecordPrefix 的 t«tid»_r 前缀），每个索引各一条规则（EncodeTableIndexPrefix 的 t«tid»_i«indexID» 前缀）。

细粒度规则的意义在于：目标表如果缺少某些索引，恢复时只需丢弃对应索引前缀的数据，而不会误伤整表数据。

正向改写的 Go 实现与文档中的 Rust 伪代码一一对应：matchOldPrefix 线性扫描找首条前缀命中的规则，replacePrefix 用 slices.Concat(rule.GetNewKeyPrefix(), s[len(rule.GetOldKeyPrefix()):]) 拼接新前缀与原 Key 的剩余部分——即文档算法的逐行直译。

Key Rewrite 对导入流程的影响

Key Rewrite 一旦引入，导入链路上就出现了两套 Key 空间：备份源文件里的 Key 是“改写前”的，而 PD 的 Region 元信息、TiKV 中 ingest 后的数据是“改写后”的。设计文档对此总结了两条不变式：

1. 源数据必然是改写前的；
2. PD 和 Ingest 后的结果必然是改写后的。

于是问题变成：在导入流程的哪个位置做改写，才能拿到最大性能？文档给出的当时导入流程为：

1. 把 KV 对写入到 RocksDB 实例（“SST File”）来排序；
2. (Prepare) 在 client 侧遍历此 SST file：
   1. (Pre-split) 找到所有 Range 的 EndKey 和 rewrite rule 的 NewPrefix，依照这些 Key 执行 BatchSplit；
   2. (GetRegion) 执行 PD 的 get_region_info 取得这些 Key 对应的 Region 信息；
   3. (Split) 用 batch_split_region 把 Region 从这些 Key 处分裂成若干部分。若出现 EpochNotMatch 或 NotLeader 错误由 Split 侧重试若干次（因为 batch 化，若交给 GetRegion 重试必须重试整个 batch），仍失败则交由 GetRegion 重试；
   4. (Scatter) 用 scatter_region 将分裂后 Range 对应的 Region 打散；
3. (Import) 按 Prepare 步获得的 Ranges 并行导入：
   1. (GetRegion) 重新取得该 Range 对应的 Region。若横跨多个 Region，则按 PD 返回的信息分割 Range 后重试；
   2. (Encode) 读取 Engine file 在这个 Range 的 KV pairs，编码成 SST 格式；
   3. (Upload) 把 SST 分别上传到各 Peers (Leader + Followers)；
   4. (Ingest) 对 Leader（或第一个 Follower）发送 ingest_sst 命令。

改写前后的 Key 在上述步骤中交叉出现：Pre-split 要用改写后的 Key（PD 只认识新 Table ID 的 Range），而 Import 读源文件时读到的还是改写前的 Key。

解决方案：在每个副本 ingest 前就地改写

文档最终选择的方案是 Key rewrite before ingest in every replica：不改写源文件、不在网络上传输改写后的数据，而是把 Rewrite Rule 随下载请求下发给 TiKV 的每个 Peer，由 TiKV 在 download（从备份存储拉取 KV 并落盘为本地 SST）时就地完成前缀替换，之后再走常规的 Raft + Ingest 路径。

用文档的成本模型评估（N 为数据 Key 数，R 为副本数）：

环节	成本	说明
源数据读盘	(R+1)·N	Split 前在 Importer 读、Key rewrite 前在 TiKV 读
Importer 写盘	0	本地引擎文件无需改写拷贝
Importer 读盘	0	—
TiKV 写盘	R·N	Key rewrite 后暂存到新 SST 文件
网络传输	0	传输的仍是原始备份数据
Key Rewrite 次数	R+1	Pre-split 一次 + 每个副本 Raft-Ingest 前一次

这个方案的本质是把改写下沉到数据所在的每个副本：Importer 侧只按改写后的 Range 做调度（Pre-split / GetRegion / Split / Scatter 全部基于 NewPrefix 派生的 Key），网络传输的内容与不做 Rewrite 时完全一致（0 额外传输），改写本身只在 TiKV 侧发生 R+1 次（一次用于计算 Pre-split 切分点，每个副本 ingest 前各一次）。

当前源码中的对应实现

当前 BR 的快照导入路径与该方案的骨架一致，只是“Encode + Upload”两步已被 TiKV 侧的 Download 接口吸收：

1. Download 阶段携带规则：Importer 调用 TiKV 的 DownloadSST 接口（import.go 中 importer.importClient.DownloadSST(dctx, peer.GetStoreId(), req)），DownloadRequest 中携带 Sst（含 Range、CRC）、RewriteRule 与 StorageBackend。TiKV 每个 Peer 从备份存储直接下载 KV、应用 Rewrite Rule 生成新前缀的本地 SST 并返回实际生效的 Range（resp.Range）供 Importer 校正 SSTMeta。从源码结构看，rewrite 的执行点正是文档所说的“每个 Peer 独自在 ingest 之前”；
2. Ingest 阶段对 Leader 发送批量 ingest：ingestSSTs 组装 MultiIngestRequest 发给 Leader，遇到 NotLeader 时从错误响应取新 Leader 重试、通过 PD GetRegion 刷新 Region 信息，EpochNotMatch 则直接报错——这与文档第 2.3 步描述的 Split 重试策略一脉相承；
3. Pre-split / Scatter 在写入之前完成：split/splitter.go 基于 rewrite 后的 Range 执行 BatchSplit 与 Scatter，保证下载时数据能落在正确且已打散的 Region 中。

规则在进入 TiKV 之前还有一道客户端侧的校验。ValidateFileRewriteRule 对每个备份文件检查三件事：

• 文件的 StartKey 能匹配到某条规则（匹配不上则报 ErrRestoreInvalidRewrite，并解出 Table ID 便于定位）；
• EndKey 同样能匹配到规则；
• StartKey 与 EndKey 命中的规则 NewKeyPrefix 必须相同——一个备份文件应只落入一个表的一个 Region，否则说明备份数据脏了或来自不兼容版本的 BR，直接报错并输出两侧规则的十六进制内容供排查。

配套的工具函数还包括：RewriteRange（把调度用的 Range 边界按规则前移，先校验起止 Key 的 Table ID 一致）、GetRewriteRawKeys（产出改写后的原始 Key 区间）以及 GetRewriteTableID（按规则把旧 Table ID 映射为新 Table ID，供日志与 checkpoint 使用）。日志层面，logutil 提供了 logutil.RewriteRule 字段，RewriteRules.String() 对前缀做了脱敏（redact）再输出，避免把真实 Key 完整打进日志。

小结

Key Rewrite 是 BR 恢复链路中连接“备份时的 Schema”与“恢复时 Schema”的枢纽机制：

• 设计层面，它用一对 bytes 前缀字段表达通用的前缀替换，摆脱了对 TiDB Key 编码格式的假设，天然支持多规则并存与空旧前缀的“统一加前缀”场景（Lightning 的用途）；
• 流程层面，改写被安排在“每个副本 ingest 前”这一位置，使 Importer 与网络侧零额外成本，改写开销收敛为 R+1 次前缀替换；
• 实现层面，从 GetRewriteRules 的规则生成、ValidateFileRewriteRule 的落盘前校验，到 import.go 中 Download/Ingest 请求携带规则下发 TiKV，2019 年讨论确定的骨架在当前 BR 代码中依然清晰可辨，并扩展出了 Keyspace 迁移与 PITR 时间过滤两项能力。

进一步阅读可参考同目录的两篇配套设计文档：新设计 与 Ingest 重组设计。