memtune

数据库内存池自适应管理框架

概述

本框架面向 GaussDB/OpenGauss，目标是在 max_process_memory 约束下，根据当前 workload 动态调整共享内存池（冷）与动态内存池（热）的分配比例，实现：

• 高并发 TP + 突发 AP 场景下 TPS 抖动控制在 10% 以内
• 低 TP 压力场景下大查询落盘降低 90% 以上

框架分为离线训练和在线推理两个阶段，最终以 C 扩展形式嵌入数据库内核，不依赖外部服务。

核心约束

• 必须内核化：调度逻辑在 GaussDB 内核中实现，不允许依赖外部挂载模型或云端协同
• 秒级响应：负载特征秒级突变时（如双十一 0 点 TPS 瞬增 10 倍、大 AP 任务启动），框架须秒级响应
• 调整无感：内存池动态调整发生在业务高峰期时，对数据库吞吐影响最小
• TP 优先：在 TP+AP 混合负载下，确保 TP 的 TPS 吞吐量平稳、latency 不抖动的前提下，再最大化 AP 性能

---

整体架构

离线阶段
  采集层 (ExecutorEnd_hook + planner_hook)
      ↓
  模型一训练：单查询内存估算 (Teacher LLM → 蒸馏 → 决策树)
      ↓
  模型二训练：池分配决策 (Teacher LLM → 蒸馏 → 决策树)
      ↓
  IR 序列化 → 存入 pg_catalog

在线阶段
  planner_hook：提取计划特征 → 模型一推理 → AP 压力估算
      ↓
  系统状态采集：TP 并发、池使用率、近期落盘率
      ↓
  模型二推理 → target_dynamic_pool_ratio
      ↓
  执行池调整 + 在线微调反馈

---

组件一：数据采集层

功能

在内核 hook 点采集训练所需的原始数据，写入日志表供离线训练使用。

实现了一个pgsql插件

Hook 点

Hook	采集内容
planner_hook	PlannedStmt 计划树特征（算子类型、rows、width）
ExecutorEnd_hook	实测峰值内存、actual rows、是否落盘

核心采集逻辑

// 遍历 PlanState 树，采集受 work_mem 管辖的算子峰值内存
static Size collect_peak_work_mem(PlanState *planstate) {
    Size total = 0;
    if (planstate == NULL) return 0;
    switch (nodeTag(planstate->plan)) {
        case T_Sort:
            total += tuplesort_get_memkb(((SortState*)planstate)->tuplesortstate) * 1024;
            break;
        case T_HashJoin:
            total += ((HashJoinState*)planstate)->hj_HashTable->spaceUsed;
            break;
        case T_Agg:
            total += MemoryContextMemAllocated(((AggState*)planstate)->hash_metacxt, true);
            break;
    }
    total += collect_peak_work_mem(planstate->lefttree);
    total += collect_peak_work_mem(planstate->righttree);
    return total;
}

日志表结构

CREATE TABLE wl_memory_samples (
    query_fingerprint      text,
    feature_vec            jsonb,       -- 计划特征
    peak_work_mem_kb       integer,     -- 标签
    actual_rows            bigint,      -- 用于计算估算误差
    estimated_rows         bigint,
    spill_occurred         boolean,
    collected_at           timestamptz DEFAULT now()
);

数据质量过滤

actual_rows / estimated_rows > 10 的样本直接丢弃，避免 planner 估算严重失准的样本污染训练集。

实现结构

数据采集层拆分为两个独立模块：

模块	语言	职责
C 扩展（mem_collector）	C	加载 hook 进 PostgreSQL，采集内存信息、算子数量等原始数据
Python 数据流水线	Python	读取 C 扩展写入的原始数据，结合离线信息加工为训练样本

C 扩展源码：https://github.com/aotenjou/mem_collector

Python 数据流水线

DB query_samples
      ↓
generate_window_samples        # 将单查询样本聚合为时间窗口快照
      ↓
window_samples
      ↓
label_window_samples           # 离线重放搜索，打 target_dynamic_pool_ratio 标签
      ↓
export_query_samples           # 导出单查询样本（模型一训练用）
      ↓
JSONL
      ↓
filter_samples                 # 过滤 estimation_error > 10 的样本
      ↓
split_dataset                  # 按时间切分训练/验证/测试集

---

组件二：模型一 —— 单查询内存估算

功能

给定一条查询的计划树特征，预测其执行所需的峰值 work_mem。输出供模型二使用，同时可直接用于单查询的 work_mem 预分配。

数据格式

输入（4 维，全部归一化）

特征	含义	归一化分母
max_sort_ratio	max(rows×width) for Sort 节点	total_server_ram
max_hash_ratio	max(build_rows×width) for HashJoin	total_server_ram
max_hashagg_ratio	max(groups×width) for HashAgg	total_server_ram
parallel_ratio	actual_workers / max_parallel_workers	—

标签（1 维）

target_mem_ratio = peak_work_mem_bytes / total_server_ram

样本示例

{
  "max_sort_ratio":    0.023,
  "max_hash_ratio":    0.047,
  "max_hashagg_ratio": 0.011,
  "parallel_ratio":    0.50,
  "target_mem_ratio":  0.061
}

AP 查询识别

窗口内含以下任一算子的查询视为 AP 查询：

bool is_ap_query(PlannedStmt *stmt) {
    return has_node_type(stmt, T_Sort)     ||
           has_node_type(stmt, T_HashJoin) ||
           has_node_type(stmt, T_Agg);
}

最重 AP 查询的选取

对窗口内每条 AP 查询计算权重分，取最高者的计划特征作为模型二的输入：

double ap_weight_score(PlannedStmt *stmt) {
    double sort_est  = max_sort_rows(stmt)      * avg_width(stmt);
    double hash_est  = max_hash_build(stmt)     * avg_width(stmt);
    double hagg_est  = max_hashagg_groups(stmt) * avg_width(stmt);
    return Max(sort_est, Max(hash_est, hagg_est));
}

取 max 而非 sum，准确反映单次内存申请的峰值压力。

---

组件三：模型二 —— 池分配决策

功能

给定当前系统状态快照，输出动态内存池应占 max_process_memory 的最优比例。

数据格式

输入（6 维）

特征	含义
tp_concurrency_ratio	当前并发 TP 查询数 / max_connections
ap_count	当前窗口内 AP 查询数量
ap_max_predicted_mem_ratio	模型一输出 × 1.2（保守系数）
dynamic_pool_current_ratio	当前动态池大小 / max_process_memory
dynamic_pool_usage_ratio	动态池使用量 / 动态池大小
spill_rate_recent	近期窗口落盘次数 / ap_count

保守系数 1.2 的作用：宁可高估 AP 内存需求，避免落盘（题目对落盘的惩罚重于对共享池偏小的惩罚）。

标签（1 维）

target_dynamic_pool_ratio ∈ [0.1, 0.9]

通过离线重放历史 workload，穷举比例步长 0.05，找到满足以下条件的最小值打标：

1. spill_count = 0
2. tps_jitter < 10%
3. dynamic_pool_usage_ratio < 0.95

样本示例

{
  "tp_concurrency_ratio":       0.72,
  "ap_count":                   3,
  "ap_max_predicted_mem_ratio": 0.073,
  "dynamic_pool_current_ratio": 0.40,
  "dynamic_pool_usage_ratio":   0.91,
  "spill_rate_recent":          0.33,
  "target_dynamic_pool_ratio":  0.55
}

---

组件四：知识蒸馏与 IR

流程

LLM (Teacher)
  ├── 扩充训练标签（合成样本）
  ├── SHAP 特征归因（验证特征设计）
  └── 生成 soft label 供蒸馏

决策树 (Student)
  ├── 模型一：深度 ≤ 4，~300 bytes
  └── 模型二：深度 ≤ 4，~300 bytes

IR 格式

每个节点序列化为：

typedef struct IRNode {
    uint8_t  feature_idx;   // 1 byte
    float    threshold;     // 4 bytes
    float    leaf_value;    // 4 bytes（仅叶节点有效）
    int16_t  left_child;    // 2 bytes（-1 表示叶节点）
    int16_t  right_child;   // 2 bytes
} IRNode;                   // 共 13 bytes/节点

两个模型合计约 600 bytes，存入系统表一行 bytea 字段。

系统表

CREATE TABLE pg_memory_model (
    model_id      integer PRIMARY KEY,
    model_name    text,            -- 'single_query' | 'pool_allocation'
    version       integer,
    model_blob    bytea,           -- 序列化 IR
    feature_schema jsonb,          -- 特征名称与顺序
    created_at    timestamptz
);

---

组件五：内核推理引擎

功能

在 PostgreSQL 内核 hook 点调用两个模型，完成实时推理并执行池调整。

推理流程

planner_hook
  ├── 识别 AP 查询，计算权重分
  ├── 提取最重 AP 查询的 4 维计划特征
  └── 调用模型一 → ap_max_predicted_mem_ratio（× 1.2）

系统状态采集
  ├── pg_stat_activity → tp_concurrency_ratio
  ├── 内存上下文 API → pool usage ratios
  └── 滑动窗口计数器 → spill_rate_recent

模型二推理
  └── target_dynamic_pool_ratio

执行调整
  └── 更新 dynamic_pool_size = ratio × max_process_memory

推理函数（C 实现）

float ir_predict(IRNode *tree, float *features) {
    int idx = 0;
    while (tree[idx].left_child != -1) {
        if (features[tree[idx].feature_idx] <= tree[idx].threshold)
            idx = tree[idx].left_child;
        else
            idx = tree[idx].right_child;
    }
    return tree[idx].leaf_value;
}
// 推理延迟：~1μs（决策树深度 4）

在线微调

每个时间窗口结束后，根据实际反馈调整模型二叶节点输出值：

// 观测到落盘 → 当前 ratio 偏低，小步上调
if (spill_occurred)
    current_leaf->leaf_value += 0.02;

// 观测到 TPS 抖动 > 10% → ratio 偏高，小步下调
if (tps_jitter > 0.10)
    current_leaf->leaf_value -= 0.02;

无需重新训练，仅修改叶节点值（几十字节写操作）。

---

组件六：泛化性保障

跨实例泛化

所有特征和标签均以 total_server_ram 或无量纲比值表示，消除绝对内存规模差异。

额外加入 3 个实例画像特征（仅模型二）：

特征	含义
memory_tier	离散化实例规模：0=≤16GB, 1=32~64GB, 2=≥128GB
storage_spill_cost	归一化落盘代价（iops_available / iops_baseline）
workload_ap_ratio	历史 AP 查询占总查询比例（滑动平均）

跨负载泛化

• 静态：实例画像特征区分不同负载模式
• 动态：在线微调机制处理分布漂移
• 兜底：估算误差过大时回退到默认 work_mem

---

数据集组织

dataset/
├── query_level/          # 模型一
│   ├── train.jsonl
│   ├── val.jsonl
│   └── test.jsonl        # 按时间切分，不随机切分（防时序泄露）
└── window_level/         # 模型二
    ├── train.jsonl
    ├── val.jsonl
    └── test.jsonl

---

关键设计决策汇总

决策	选择	原因
标签粒度	单一连续标量	直接对应调控目标，蒸馏损失最小
归一化分母	total_server_ram	稳定，与内存分配决策直接相关
SQL 解析	不解析原始 SQL	计划树已融合 SQL 结构与数据分布
模型结构	双模型串联	任务分解，各自模型更浅、精度更高
IR 格式	序列化决策树	~300 bytes/模型，1μs 推理，50 行 C 代码
泛化方案	归一化 + 实例画像 + 在线微调	分别应对静态差异和动态漂移