TiKV 内存管理优化：支持自定义消息拒绝逻辑

背景介绍

在分布式存储系统TiKV中，内存管理是一个至关重要的环节。当系统内存资源紧张时，TiKV需要采取适当的措施来防止内存耗尽导致的服务不可用。目前TiKV已经实现了基于内存使用情况的raft消息拒绝机制，但这种机制存在一定的局限性。

现有机制分析

当前TiKV的内存管理机制主要针对raft消息处理，具体规则如下：

1. 当系统内存使用达到高水位线(high water)时触发保护机制
2. 当raft集合消耗的内存超过总使用内存的reject_messages_on_memory_ratio比例时，会拒绝新的msgAppend消息

这种机制在基于RocksDB引擎的TiKV中工作良好，但对于使用不同存储引擎的场景(如TiFlash的列式存储引擎)则存在不足。

问题挑战

在TiFlash等列式存储引擎场景下，现有机制面临两个主要挑战：

1. 快照处理内存消耗：列式存储引擎在应用快照时需要消耗大量内存，但当前机制无法针对msgSnapshot消息进行限制
2. 内存分配差异：列式存储引擎的内存使用模式与传统行式存储不同，大部分内存消耗发生在存储引擎层而非raft层

技术方案

为了解决这些问题，TiKV需要支持更灵活的内存管理策略：

1. 扩展消息类型支持：不仅限于msgAppend，还需要支持对msgSnapshot等消息类型的拒绝逻辑
2. 引入存储引擎感知：内存管理需要了解不同存储引擎的内存使用特性
3. 自定义拒绝策略：允许不同存储引擎实现自己的内存评估和拒绝逻辑

实现思路

具体实现上可以考虑以下技术路线：

1. 抽象接口设计：定义内存管理接口，允许存储引擎实现自定义逻辑
2. 内存使用反馈机制：存储引擎定期报告其内存使用情况
3. 动态策略调整：根据当前工作负载和存储引擎类型动态调整拒绝策略
4. 分级保护机制：针对不同消息类型设置不同的保护阈值

预期收益

这种改进将为系统带来以下好处：

1. 更好的资源隔离：防止单一组件耗尽系统内存
2. 更高的系统稳定性：在内存压力下更优雅地降级
3. 更灵活的架构：支持不同类型的存储引擎协同工作
4. 更精细的控制：可以根据实际业务需求调整保护策略

总结

TiKV的内存管理机制需要从简单的固定策略演进为可插拔的灵活架构。通过支持自定义的消息拒绝逻辑，TiKV能够更好地适应不同存储引擎的特性和各种工作负载场景，为上层应用提供更稳定可靠的服务。这种改进也体现了TiKV作为云原生数据库核心组件的设计理念：灵活、可扩展和自适应。