踩坑rust的partial copy導(dǎo)致metrics丟失

然而，在使用過程中，我們遇到了上報(bào)的 metrics 全部為 0 的問題。

1最小復(fù)現(xiàn)

由于使用了try_stream宏來生成 stream，因此我們懷疑在try_stream生成的代碼中有 bug 導(dǎo)致 metrics 丟失。于是我們用 cargo-expand [3] 來將查看宏生成的代碼。展開后的代碼如下：

fninto_stream_inner(
mutself,
)->implStream
>{
::from_generator(staticmove|
mut__task_context:::ResumeTy,
|->::Result<(),?StorageError>{
let():()={
letinner=self.inner;
letmutinner=inner;
#[allow(unused_mut)]
letmutinner=unsafe{
::new_unchecked(
&mutinner,
)
};
whileletSome((key,value))
={
letmut__pinned=inner.try_next();
loop{
iflet::Ready(
result,
)=unsafe{
poll(Pin::as_mut(&mut__pinned),get_context(__task_context))
}{
breakresult;
}
__task_context=(yield::Pending);
}
}?
{
self.stats.total_items+=1;
self.stats.total_size+=key.encoded_len()+value.len();
__task_context=(yield::Ready((
key,
value,
)));
}
};
#[allow(unreachable_code)]
{
return::Ok(());
loop{
__task_context=(yield::Pending);
}
}
})
}

可以看到，try_stream宏生成的代碼中，包含了一個(gè) rust generator 的閉包。閉包中收集和上報(bào) metrics 的邏輯與原代碼基本相同，按照我們對 rust 的理解，仍然不應(yīng)該會出現(xiàn) metrics 丟失的問題。因此我們懷疑是 rust 編譯器中與 generator 相關(guān)的邏輯存在問題。在 rust playground 上，我們嘗試了以下代碼來對問題進(jìn)行復(fù)現(xiàn)。

structStat{
count:usize,
vec:Vec,
}

implDropforStat{
fndrop(&mutself){
println!("count:{}",self.count);
}
}

fnmain(){

letmutstat=Stat{
count:0,
vec:Vec::new(),
};

letmutf=move||{
stat.count+=1;
1
};

println!("num:{}",f());
}

執(zhí)行以后輸出如下。

num:1
count:0

按照預(yù)期，輸出的 num 和 count 應(yīng)該都為1，因?yàn)樵谡{(diào)用閉包 f 時(shí)stat.count += 1被調(diào)用了，但是以上代碼中遇到了和最開始同樣的問題。因此以上代碼可以作為我們問題的一個(gè)最小復(fù)現(xiàn)。

2問題分析

對以上代碼進(jìn)行分析的話，我們看到閉包 f 的代碼中使用了 move，因此在閉包中使用過的對象的 ownership 應(yīng)該都會轉(zhuǎn)移到閉包中。而 struct Stats實(shí)現(xiàn)了Drop，因此Stats是不可以 partial move 的，其必須作為一個(gè)整體被 move 進(jìn)入閉包。而在閉包中執(zhí)行了stats.count += 1，因此 stats 中的 count 應(yīng)該被置為1。但是從程序的輸出可以看到在 stats 被 drop 時(shí)，stats 的 count 是 0。

我們嘗試將閉包 f 改為如下代碼，來顯式的將 stats 的 ownership 給 move 進(jìn)閉包里。

letmutf=move||{
letmutstat=stat;
stat.count+=1;
1
};

輸出恢復(fù)正常。

num:1
count:1

我們再次嘗試在閉包 f 中使用 stat 中的另一個(gè)字段 vec：

letmutf=move||{
let_=stat.vec.len();
stat.count+=1;
1
};

輸出同樣恢復(fù)正常。

num:1
count:1

可以看到，我們顯式地將 stat 整個(gè) move 進(jìn)閉包，或者在閉包中使用類型為 vec 的字段，都會使得 stat 的ownership 被 move 進(jìn)閉包。

于是我們推測，盡管 stat 實(shí)現(xiàn)了自己的 drop 導(dǎo)致不能被 partial move，但是如果我們在 move 的閉包中只使用了 stat 中實(shí)現(xiàn)了 Copy 類型的字段，則這個(gè)字段的值會被 Copy 到閉包中，而閉包中對這個(gè)字段的修改只會作用于被 Copy 后的值，而原字段并不會改變。

3驗(yàn)證猜想

我們可以通過將以上代碼編譯成二進(jìn)制代碼后，對其匯編代碼進(jìn)行分析，從而驗(yàn)證我們的猜想。然而，編譯后的匯編代碼會過于復(fù)雜且晦澀難懂，同時(shí)編譯器對其進(jìn)行的一些優(yōu)化也會改變原有的邏輯導(dǎo)致匯編代碼難以理解。因此我們打算通過分析在編譯過程中產(chǎn)生的 MIR 中間代碼來對問題進(jìn)行分析。在 rust playground 上可以十分方便地生成 MIR 代碼。

首先我們對存在問題的最小復(fù)現(xiàn)代碼生成 MIR，生成后與閉包相關(guān)的 MIR 如下?？梢钥吹竭@個(gè)閉包確實(shí)只包含了一個(gè)類型為 usize 的字段，這個(gè)字段的值取的是 stat 中的 count 值。

bb1:{
_1=Stat{count:const0_usize,vec:move_2};
_3={closure@src/main.rs:19:17:19:24}{stat:(_1.0:usize)};
}

而我們對后續(xù)測試中有正常輸出的代碼生成 MIR，生成后與閉包相關(guān)的 MIR 如下。可以看到這個(gè)閉包將整個(gè) stat 的 ownership 給 move 了進(jìn)去。

bb1:{
_1=Stat{count:const0_usize,vec:move_2};
_3={closure@src/main.rs:19:17:19:24}{stat:move_1};
}

于是，我們的猜想得到了驗(yàn)證，在我們出現(xiàn)問題的代碼中，閉包確實(shí)沒有捕獲 stat 的 ownership。

4后續(xù)與總結(jié)

我們向 rust 社區(qū)反映了這個(gè)問題，得到的反饋是，這個(gè)是 rust 2021 后實(shí)現(xiàn)的一個(gè) feature [4]。在 rust 2021 中，一個(gè)使用了 move 的閉包在捕獲一個(gè) struct 的時(shí)候，會盡可能少地去捕獲 struct 中的字段。

如果一個(gè) struct 沒有實(shí)現(xiàn) Drop，這意味著他里面的字段可以被分開 move，而閉包只會捕獲閉包中用到的字段。

如果某個(gè)被閉包使用的字段實(shí)現(xiàn)了 Copy，那他閉包并不會捕獲這個(gè)字段的 ownership，而是將這個(gè)字段 copy 一份放在閉包中。

如果一個(gè) struct 實(shí)現(xiàn)了 Drop，那他里面的字段只能作為一個(gè)整體被捕獲。但如果閉包只使用了這個(gè)閉包中實(shí)現(xiàn)了 Copy 的字段，那這個(gè)閉包不會捕獲這個(gè) struct，而是將使用到的字段 copy 一份。

我們的代碼中，正是因?yàn)檫@個(gè)行為，導(dǎo)致我們的代碼產(chǎn)生了歧義，而出現(xiàn)了 metrics 的丟失。

針對這個(gè)問題，我們認(rèn)為有兩個(gè)地方有提升的空間。

首先，try_stream這個(gè)宏的封裝存在一定的問題。在使用宏來聲明代碼中，其暴露出來的使用方法是通過調(diào)用一個(gè)方法來生成 stream，而在調(diào)用方法時(shí)，如果參數(shù)是通過 move ownership 的形式傳入的，同時(shí)在生成 stream 的代碼中我們使用了這個(gè)參數(shù)，那我們應(yīng)該認(rèn)為這個(gè) stream 包含了這個(gè)參數(shù)的 ownership。然而，由于這個(gè)宏在實(shí)現(xiàn)的時(shí)候使用了閉包，導(dǎo)致這個(gè) stream 并沒有包含這個(gè)參數(shù)的 ownership，從而導(dǎo)致問題。這個(gè)是宏封裝邏輯時(shí)的問題。

其次，rust 在語言設(shè)計(jì)上，由于引入了這個(gè)閉包捕獲 ownership 的特殊邏輯，導(dǎo)致會寫出有歧義的代碼。例如，在上述代碼中，很難想象stat.count += 1并沒有去修改 stat 中的 count 值。我們也向 rust 社區(qū)反映了這個(gè)問題 [5]。

最后，回到我們最開始的問題中。要想解決 metrics 丟失的問題，在我們的代碼中，我們只需要做以下修改就能讓代碼正常運(yùn)行[6]。

#[try_stream(ok=StateStoreIterItem,error=StorageError)]
asyncfninto_stream_inner(mutself){
letinner=self.inner;
...
self.stats.total_items+=1;
self.stats.total_size+=key.encoded_len()+value.len();
}

修改為

#[try_stream(ok=StateStoreIterItem,error=StorageError)]
asyncfninto_stream_inner(self){
letinner=self.inner;
letmutstats=self.stats;
...
stats.total_items+=1;
stats.total_size+=key.encoded_len()+value.len();
}

審核編輯：湯梓紅