更快的tsv解析代碼分享

最近在B站沖浪時(shí)發(fā)現(xiàn)一個(gè) Rust 和 Go 解析 tsv 文件的視頻, 作者需要解析使用get-NetTCPConnection | Format-Table -Property LocalAddress,LocalPort,RemoteAddress,RemotePort,State,OwningProcess獲取的本地所有 TCP 連接信息, 文件輸出大致如下

LocalAddressLocalPortRemoteAddressRemotePortStateOwningProcess
--------------------------------------------------------------
192.168.1.454339104.210.1.98443Established4504

視頻作者使用 regex 正則庫(kù)處理輸出, 發(fā)現(xiàn)比 Go 版本慢, 優(yōu)化后雖然比 Go 快, 但并沒(méi)有領(lǐng)先多少, 于是我自己嘗試使用別的優(yōu)化方法, 解析耗時(shí)能優(yōu)化使用正則解析的 10% 左右. 下面來(lái)看看我的優(yōu)化過(guò)程.

?更快的 tsv 解析[1]

?項(xiàng)目搭建[2]

?regex 解析[3]

?減少內(nèi)存分配[4]

?使用 ascii 正則[5]

?拋棄 regex[6]

?手寫(xiě)解析狀態(tài)機(jī)[7]

?SIMD 加速?[8]

?總結(jié)[9]

項(xiàng)目搭建

進(jìn)行性能時(shí)建議使用criterion[10], 它幫我們解決了性能的內(nèi)存預(yù)加載, 操作耗時(shí), 性能記錄, 圖表輸出等功能.

cargonew--libtsv
cdtsv
cargoaddcriterion--dev-Fhtml_reports
cargoaddregex

然后在 Cargo.toml 里添加如下bench 文件

[[bench]]
name="parse"
harness=false

//benches/parse.rs
#![allow(dead_code)]
usecriterion::{black_box,criterion_group,criterion_main,Criterion};

constOUTPUT:&str=include_str!("net.tsv");

fncriterion_benchmark(c:&mutCriterion){
todo!()
}

criterion_group!(benches,criterion_benchmark);
criterion_main!(benches);

測(cè)試使用的 tsv 一共 380 行.

regex 解析

使用正則解析的正則表達(dá)式很簡(jiǎn)單, 這里直接給代碼, 為了避免重復(fù)編譯正則表達(dá)式和重新分配內(nèi)存報(bào)錯(cuò)結(jié)果列表, 這里將她們作為參數(shù)傳給解析函數(shù).

structOwnedRecord{
local_addr:String,
local_port:u16,
remote_addr:String,
remote_port:u16,
state:String,
pid:u64,
}
fnregex_owned(input:&str,re:®ex::Regex,result:&mutVec){
input.lines().for_each(|line|{
ifletSome(item)=re.captures(line).and_then(|captures|{
let(_,[local_addr,local_port,remote_addr,remote_port,state,pid])=
captures.extract();
letret=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 文件里要加上使用的正則和提前創(chuàng)建好列表, 并且將函數(shù)添加的 bench 目標(biāo)里

fncriterion_benchmark(c:&mutCriterion){
letre=regex::new(r"(S+)s+(d+)s+(S+)s+(d+)s+(S+)s+(d+)").unwrap();
letmutr1=Vec::with_capacity(400);
c.bench_function("regex_owned",|b|{
b.iter(||{
//重置輸出vector
r1.clear();
regex_owned(black_box(OUTPUT),&re,&mutr1);
})
});
}

接著跑cargo bench --bench parse進(jìn)行測(cè)試, 在我的電腦上測(cè)得每次運(yùn)行耗時(shí) 450 μs 左右.

減少內(nèi)存分配

一個(gè)最簡(jiǎn)單的優(yōu)化是使用&str以減少每次創(chuàng)建String帶來(lái)的內(nèi)存分配和數(shù)據(jù)復(fù)制.

structRecord<'a>{
local_addr:&'astr,
local_port:u16,
remote_addr:&'astr,
remote_port:u16,
state:&'astr,
pid:u64,
}

兩個(gè)函數(shù)代碼差不多, 所以這里不再列出來(lái), 可以通過(guò)gits: tsv 解析[11]獲取完整代碼.

可惜這次改動(dòng)帶來(lái)的優(yōu)化非常小, 在我的電腦上反復(fù)測(cè)量, 這個(gè)版本耗時(shí)在 440 μs 左右.

使用 ascii 正則

rust 的 regex 正則默認(rèn)使用 unicode, 相比于 ascii 編碼, unicode 更復(fù)雜, 因此性能也相對(duì)較低, 剛好要解析的內(nèi)容都是ascii字符, 使用 ascii 正則是否能提升解析速度呢? regex 有regex::bytes模塊用于 ascii 解析, 但為了適配字段, 這里不得不使用transmute將&[u8]強(qiáng)制轉(zhuǎn)換成&str

fncast(data:&[u8])->&str{
unsafe{std::transmute(data)}
}
fnregex_ascii<'a>(input:&'astr,re:®ex::Regex,result:&mutVec>){
input.lines().for_each(|line|{
ifletSome(item)=re.captures(line.as_bytes()).and_then(|captures|{
let(_,[local_addr,local_port,remote_addr,remote_port,state,pid])=
captures.extract();
letret=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

鑒于內(nèi)容格式比較簡(jiǎn)單, 如果只使用 rust 內(nèi)置的 split 等方法解析性能會(huì)不會(huì)更好呢? 解析思路很簡(jiǎn)單, 使用lines得到一個(gè)逐行迭代器, 然后對(duì)每行使用 split 切分空格再逐個(gè)解析即可

fnsplit<'a>(input:&'astr,result:&mutVec>){
input
.lines()
.filter_map(|line|{
letmutiter=line.split(['','	','
']).filter(|c|!c.is_empty());
letlocal_addr=iter.next()?;
letlocal_port:u16=iter.next()?.parse().ok()?;
letremote_addr=iter.next()?;
letremote_port:u16=iter.next()?.parse().ok()?;
letstate=iter.next()?;
letpid: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只后還需要過(guò)濾不是空白的字符串, 這是因?yàn)樽址?a b"split 之后得到["a", "", "b"].

經(jīng)測(cè)試, 這個(gè)版本測(cè)試耗時(shí)大概為 53 μs, 這真是一個(gè)巨大提升, rust 的 regex 性能確實(shí)有些問(wèn)題.

每次 split 之后還需要 filter 感覺(jué)有些拖沓, 剛好有個(gè)split_whitespace[12], 換用這個(gè)方法, 將新的解析方法命名為split_whitespace后再測(cè)試下性能

letmutiter=line.split_whitespace();

令人意想不到的是性能居然倒退了, 這次耗時(shí)大概 60 μs, 仔細(xì)研究下來(lái)還是 unicode 的問(wèn)題, 改用 ascii 版本的split_ascii_whitespace之后性能提升到 45 μs.

手寫(xiě)解析狀態(tài)機(jī)

除了上述的方法, 我還嘗試將 Record 的 local_addr 和 remote_addr 改成std::IpAddr, 消除next()?.parse().ok()?等其他方法, 但收益幾乎沒(méi)有, 唯一有作用的辦法是手寫(xiě)解析狀態(tài)機(jī).

大致思路是, 對(duì)于輸出來(lái)說(shuō), 我們只關(guān)系它是以下三種情況

1.換行符 NL

2.除了換行符的空白符 WS

3.非空白字符 CH

只解析 LocalAddr 和 LocalPort 解析狀態(tài)機(jī)如下, 如果要解析更多字段, 按順序添加即可.

因?yàn)榇a有些復(fù)雜, 所以這里不再貼出來(lái), 完整代碼在 gits 上. 手寫(xiě)狀態(tài)機(jī)的版本耗時(shí)大概在 32 μs 左右. 這版本主要性能提升來(lái)自手寫(xiě)狀態(tài)機(jī)減少了循環(huán)內(nèi)的分支判斷.

SIMD 加速?

在上面手寫(xiě)解析的例子里, 處理過(guò)程類(lèi)似與將輸出作為一個(gè) vec, 狀態(tài)機(jī)作為另一個(gè) vec, 將兩個(gè) vec 進(jìn)行某種運(yùn)算后輸出結(jié)果, 應(yīng)該能使用 simd 進(jìn)行加速, 但我還沒(méi)想出高效實(shí)現(xiàn). 所以這里只給出可能的參考資料

1.zsv[13]使用 simd 加速的 csv 解析庫(kù)

2.simd base64[14]一篇介紹使用 simd 加速 base64 解析的博客, 非常推薦

總結(jié)

rust regex 在某時(shí)候確實(shí)存在性能問(wèn)題, 有時(shí)候使用簡(jiǎn)單的 split 的方法手動(dòng)解析反而更簡(jiǎn)單性能也更高, 如果情況允許, 使用 ascii 版本能進(jìn)一步提升性能, 如果你追求更好的性能, 手寫(xiě)一個(gè)狀態(tài)不失為一種選擇, 當(dāng)然我不建議在生產(chǎn)上這么做. 同時(shí)我也期待有 simd 加速的例子.

審核編輯：黃飛