fix(utils): TaskGuard::drop 不再调用 tokio::task::id()
问题:TaskGuard::drop 中调用 tokio::task::id()。abort 路径下 future 的销毁若发生在非任务上下文,id() 直接 panic 刷日志; 若发生在调用方任务上下文,则拿到调用方的任务 id,把它从 任务组中误删,破坏任务组记账。 修复:spawn 时在任务上下文内获取自身 id 存入 TaskGuard, Drop 只使用保存的 id。 测试:test_abort_task_keeps_caller_bookkeeping 验证在同组任务内 abort 其他任务后,调用方条目不被误删、wait_all_stopped 正常返回。
This commit is contained in:
@@ -36,12 +36,16 @@ impl TaskGroupInner {
|
|||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
|
||||||
let guard = TaskGuard {
|
let weak = Arc::downgrade(self);
|
||||||
inner: Arc::downgrade(self),
|
|
||||||
};
|
|
||||||
|
|
||||||
let handle = tokio::spawn(async move {
|
let handle = tokio::spawn(async move {
|
||||||
let _guard = guard;
|
// 在任务上下文内获取自身 id 存入 guard;
|
||||||
|
// 不能延迟到 Drop 里调 tokio::task::id():
|
||||||
|
// abort 路径下 future 可能在非任务上下文被销毁(panic),
|
||||||
|
// 在调用方任务上下文被销毁时又会拿到错误的 id 误删条目
|
||||||
|
let _guard = TaskGuard {
|
||||||
|
inner: weak,
|
||||||
|
task_id: tokio::task::id(),
|
||||||
|
};
|
||||||
f.await;
|
f.await;
|
||||||
});
|
});
|
||||||
|
|
||||||
@@ -107,13 +111,14 @@ impl Drop for TaskGroupInner {
|
|||||||
|
|
||||||
struct TaskGuard {
|
struct TaskGuard {
|
||||||
inner: Weak<TaskGroupInner>,
|
inner: Weak<TaskGroupInner>,
|
||||||
|
/// 创建时(任务上下文内)获取的自身任务 id
|
||||||
|
task_id: Id,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Drop for TaskGuard {
|
impl Drop for TaskGuard {
|
||||||
fn drop(&mut self) {
|
fn drop(&mut self) {
|
||||||
if let Some(inner) = self.inner.upgrade() {
|
if let Some(inner) = self.inner.upgrade() {
|
||||||
let task_id = tokio::task::id();
|
inner.remove_task(self.task_id);
|
||||||
inner.remove_task(task_id);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -286,4 +291,41 @@ mod tests {
|
|||||||
.expect("wait_all_stopped should return after all tasks complete")
|
.expect("wait_all_stopped should return after all tasks complete")
|
||||||
.unwrap();
|
.unwrap();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// abort 路径:任务被 stop() 终止后,TaskGuard 必须用创建时保存的 id
|
||||||
|
/// 注销自身;若在 Drop 里调 tokio::task::id(),在非任务上下文会 panic,
|
||||||
|
/// 在调用方任务上下文则会误删调用方的条目。
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_abort_task_keeps_caller_bookkeeping() {
|
||||||
|
let manager = TaskGroupManager::new();
|
||||||
|
let (group, _guard) = manager.create_task().unwrap();
|
||||||
|
|
||||||
|
let victim = group.spawn(async {
|
||||||
|
tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
|
||||||
|
});
|
||||||
|
|
||||||
|
// observer 在同组任务内 abort victim,随后挂起等待放行
|
||||||
|
let (done_tx, done_rx) = tokio::sync::oneshot::channel::<()>();
|
||||||
|
let (exit_tx, exit_rx) = tokio::sync::oneshot::channel::<()>();
|
||||||
|
let observer = group.spawn(async move {
|
||||||
|
victim.stop().await;
|
||||||
|
let _ = done_tx.send(());
|
||||||
|
let _ = exit_rx.await;
|
||||||
|
});
|
||||||
|
|
||||||
|
done_rx.await.unwrap();
|
||||||
|
// victim 的 guard 注销不得误删 observer 的条目
|
||||||
|
assert!(
|
||||||
|
observer.is_running(),
|
||||||
|
"aborting victim must not remove the caller's task entry"
|
||||||
|
);
|
||||||
|
|
||||||
|
let _ = exit_tx.send(());
|
||||||
|
tokio::time::timeout(
|
||||||
|
std::time::Duration::from_secs(2),
|
||||||
|
group.wait_all_stopped(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("wait_all_stopped should return after observer exits");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user