更快的 tsv 解析
最近在B站冲浪时发现一个 Rust 和 Go 解析 tsv 文件的视频, 作者需要解析使用 get-NetTCPConnection | Format-Table -Property LocalAddress,LocalPort,RemoteAddress,RemotePort,State,OwningProcess
获取的本地所有 TCP 连接信息, 文件输出大致如下
LocalAddress LocalPort RemoteAddress RemotePort State OwningProcess ------------ --------- ------------- ---------- ----- ------------- 192.168.1.4 54339 104.210.1.98 443 Established 4504
视频作者使用 regex 正则库处理输出, 发现比 Go 版本慢, 优化后虽然比 Go 快, 但并没有领先多少, 于是我自己尝试使用别的优化方法, 解析耗时能优化使用正则解析的 10% 左右. 下面来看看我的优化过程.
项目搭建
进行性能时建议使用 criterion, 它帮我们解决了性能的内存预加载, 操作耗时, 性能记录, 图表输出等功能.
cargo new --lib tsv cd tsv cargo add criterion --dev -F html_reports cargo add regex
然后在 Cargo.toml 里添加如 bench 文件
[[bench]] name = "parse" harness = false
// benches/parse.rs #![allow(dead_code)] use criterion::{black_box, criterion_group, criterion_main, Criterion}; const OUTPUT: &str = include_str!("net.tsv"); fn criterion_benchmark(c: &mut Criterion) { todo!() } criterion_group!(benches, criterion_benchmark); criterion_main!(benches);
测试使用的 tsv 一共 380 行.
regex 解析
使用正则解析的正则表达式很简单, 这里直接给代码, 为了避免重复编译正则表达式和重新分配内存报错结果列表, 这里将她们作为参数传给解析函数.
struct OwnedRecord { local_addr: String, local_port: u16, remote_addr: String, remote_port: u16, state: String, pid: u64, } fn regex_owned(input: &str, re: ®ex::Regex, result: &mut Vec<OwnedRecord>) { input.lines().for_each(|line| { if let Some(item) = re.captures(line).and_then(|captures| { let (_, [local_addr, local_port, remote_addr, remote_port, state, pid]) = captures.extract(); let ret = OwnedRecord { local_addr: local_addr.to_string(), local_port: local_port.parse().ok()?, remote_addr: remote_addr.to_string(), remote_port: remote_port.parse().ok()?, state: state.to_string(), pid: pid.parse().ok()?, }; Some(ret) }) { result.push(item); } }); assert_eq!(result.len(), 377); }
parse.rs 文件里要加上使用的正则和提前创建好列表, 并且将函数添加的 bench 目标里
fn criterion_benchmark(c: &mut Criterion) { let re = regex::Regex::new(r"(\S+)\s+(\d+)\s+(\S+)\s+(\d+)\s+(\S+)\s+(\d+)").unwrap(); let mut r1 = Vec::with_capacity(400); c.bench_function("regex_owned", |b| { b.iter(|| { // 重置输出 vector r1.clear(); regex_owned(black_box(OUTPUT), &re, &mut r1); }) }); }
接着跑 cargo bench --bench parse
进行测试, 在我的电脑上测得每次运行耗时 450 µs 左右.
减少内存分配
一个最简单的优化是使用 &str
以减少每次创建 String
带来的内存分配和数据复制.
struct Record<'a> { local_addr: &'a str, local_port: u16, remote_addr: &'a str, remote_port: u16, state: &'a str, pid: u64, }
两个函数代码差不多, 所以这里不再列出来, 可以通过 gits: tsv 解析 获取完整代码.
可惜这次改动带来的优化非常小, 在我的电脑上反复测量, 这个版本耗时在 440 µs 左右.
使用 ascii 正则
rust 的 regex 正则默认使用 unicode, 相比于 ascii 编码, unicode 更复杂, 因此性能也相对较低, 刚好要解析的内容都是ascii字符, 使用 ascii 正则是否能提升解析速度呢? regex 有 regex::bytes
模块用于 ascii 解析, 但为了适配字段, 这里不得不使用 transmute
将 &[u8]
强制转换成 &str
fn cast(data: &[u8]) -> &str { unsafe { std::mem::transmute(data) } } fn regex_ascii<'a>(input: &'a str, re: ®ex::bytes::Regex, result: &mut Vec<Record<'a>>) { input.lines().for_each(|line| { if let Some(item) = re.captures(line.as_bytes()).and_then(|captures| { let (_, [local_addr, local_port, remote_addr, remote_port, state, pid]) = captures.extract(); let ret = Record { local_addr: cast(local_addr), local_port: cast(local_port).parse().ok()?, remote_addr: cast(remote_addr), remote_port: cast(remote_port).parse().ok()?, state: cast(state), pid: cast(pid).parse().ok()?, }; Some(ret) }) { result.push(item); } }); assert_eq!(result.len(), 377); }
添加到 bench 后性能大概多少呢?, 很遗憾, 性能与 regex_borrow 差不多, 在 430 µs 左右.
抛弃 regex
鉴于内容格式比较简单, 如果只使用 rust 内置的 split 等方法解析性能会不会更好呢?
解析思路很简单, 使用 lines
得到一个逐行迭代器, 然后对每行使用 split 切分空格再逐个解析即可
fn split<'a>(input: &'a str, result: &mut Vec<Record<'a>>) { input .lines() .filter_map(|line| { let mut iter = line.split([' ', '\t', '\r']).filter(|c| !c.is_empty()); let local_addr = iter.next()?; let local_port: u16 = iter.next()?.parse().ok()?; let remote_addr = iter.next()?; let remote_port: u16 = iter.next()?.parse().ok()?; let state = iter.next()?; let pid: u64 = iter.next()?.parse().ok()?; Some(Record { local_addr, local_port, remote_addr, remote_port, state, pid, }) }) .for_each(|item| result.push(item)); assert_eq!(result.len(), 377); }
注意 line.split
只后还需要过滤不是空白的字符串, 这是因为字符串 "a b"
split 之后得到 ["a", "", "b"]
.
经测试, 这个版本测试耗时大概为 53 µs, 这真是一个巨大提升, rust 的 regex 性能确实有些问题.
每次 split 之后还需要 filter 感觉有些拖沓, 刚好有个split_whitespace, 换用这个方法, 将新的解析方法命名为split_whitespace
后再测试下性能
let mut iter = line.split_whitespace();
令人意想不到的是性能居然倒退了, 这次耗时大概 60 µs, 仔细研究下来还是 unicode 的问题, 改用 ascii 版本的 split_ascii_whitespace
之后性能提升到 45 µs.
手写解析状态机
除了上述的方法, 我还尝试将 Record 的 local_addr 和 remote_addr 改成 std::net::IpAddr
, 消除 next()?.parse().ok()?
等其他方法, 但收益几乎没有, 唯一有作用的办法是手写解析状态机.
大致思路是, 对于输出来说, 我们只关系它是以下三种情况
- 换行符 NL
- 除了换行符的空白符 WS
- 非空白字符 CH
只解析 LocalAddr 和 LocalPort 解析状态机如下, 如果要解析更多字段, 按顺序添加即可.
stateDiagram-v2 [*] --> LineStart LineStart --> LineStart:WS LineStart --> ConsumeLocalAddr:Ch ConsumeLocalAddr --> LineStart:NL ConsumeLocalAddr --> LocalAddr:WS ConsumeLocalAddr --> ConsumeLocalAddr:CH LocalAddr --> LineStart:NL LocalAddr --> LocalAddr:WS LocalAddr --> ConsumeLocalPort:CH ConsumeLocalPort --> ParsePort:NL ConsumeLocalPort --> ParsePort:WS ConsumeLocalPort --> ConsumeLocalPort:CH ParsePort --> LineDone:parse ok ParsePort --> Abort:parse error LineDone --> LineStart:NL Abort --> LineStart:NL LineDone --> [*]:EOF Abort --> [*]:EOF
因为代码有些复杂, 所以这里不再贴出来, 完整代码在 gits 上. 手写状态机的版本耗时大概在 32 µs 左右. 这版本主要性能提升来自手写状态机减少了循环内的分支判断.
SIMD 加速?
在上面手写解析的例子里, 处理过程类似与将输出作为一个 vec, 状态机作为另一个 vec, 将两个 vec 进行某种运算后输出结果, 应该能使用 simd 进行加速, 但我还没想出高效实现. 所以这里只给出可能的参考资料
- zsv 使用 simd 加速的 csv 解析库
- simd base64 一篇介绍使用 simd 加速 base64 解析的博客, 非常推荐
总结
rust regex 在某时候确实存在性能问题, 有时候使用简单的 split 的方法手动解析反而更简单性能也更高, 如果情况允许, 使用 ascii 版本能进一步提升性能, 如果你追求更好的性能, 手写一个状态不失为一种选择, 当然我不建议在生产上这么做. 同时我也期待有 simd 加速的例子.