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

背景介绍

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

现有机制分析

当前TiKV的内存拒绝机制遵循以下规则：

1. 当内存使用达到高水位线时触发限制
2. 当raft集合消耗的内存超过总使用内存的reject_messages_on_memory_ratio比例时，拒绝新的msgAppend消息

这种机制主要针对传统的基于RocksDB的存储引擎设计，能够有效防止raft日志过度消耗内存。然而，随着TiKV支持更多类型的存储引擎，特别是TiFlash这样的列式存储引擎，现有机制显示出了一些不足。

问题挑战

当使用TiFlash等列式存储引擎时，会遇到以下特殊的内存使用模式：

1. 快照应用内存消耗：列式存储引擎在应用快照时会消耗大量内存，需要能够拒绝msgSnapshot消息
2. 内存分布变化：列式存储引擎的内存主要消耗在自身的存储结构中，而非raft集合中
3. 内存管理粒度：需要更细粒度的内存控制，能够区分不同类型消息的内存影响

技术方案

为了解决这些问题，TiKV需要实现一套可扩展的内存管理框架，主要包括以下改进：

1. 消息类型感知

系统需要能够识别不同类型的raft消息（如msgAppend、msgSnapshot等），并为每种类型配置独立的内存控制策略。这包括：

• 为每种消息类型设置独立的内存使用阈值
• 实现基于消息类型的优先级控制
• 支持动态调整各类型消息的内存配额

2. 存储引擎集成接口

设计统一的存储引擎内存报告接口，允许：

• 存储引擎上报其内存使用情况
• 系统获取整体内存视图
• 基于多维度数据做出内存控制决策

3. 动态策略调整

实现可动态配置的内存管理策略，包括：

• 运行时调整内存限制参数
• 支持基于工作负载的自适应调整
• 提供详细的监控指标用于策略优化

实现考虑

在具体实现上，需要注意以下技术细节：

1. 性能影响：内存检查逻辑需要高效，避免成为性能瓶颈
2. 一致性保证：消息拒绝不能影响raft协议的最终一致性
3. 资源隔离：不同租户/业务的内存使用需要适当隔离
4. 故障恢复：内存压力缓解后需要能够快速恢复正常服务

预期收益

通过这项改进，TiKV将获得以下优势：

1. 更好的异构存储支持：能够更有效地支持TiFlash等列式存储引擎
2. 更精细的内存控制：根据实际使用模式优化内存分配
3. 更高的系统稳定性：防止特定场景下的内存耗尽
4. 更灵活的部署选项：适应不同硬件配置和工作负载

总结

TiKV的内存管理机制正在向更灵活、更智能的方向演进。通过支持自定义的消息拒绝逻辑，TiKV能够更好地适应多样化的存储引擎和不断变化的工作负载需求。这项改进不仅解决了当前与TiFlash集成时的问题，也为未来支持更多类型的存储引擎奠定了基础，是TiKV作为通用分布式存储系统的重要演进方向。