SpawnReady
主线位置:就绪主线。它只改
poll_ready:内层没就绪时,把「等就绪」这件事从调用方的任务挪到一个新任务里。请求主线上只是把错误转成BoxError。
flowchart TB
pr["poll_ready"] --> st{Inner 是哪个}
st -->|"Service(Some)"| ask[在本任务里问内层]
ask -->|Ready| ret[返回 Ready]
ask -->|Pending| mv["取出服务,spawn ready_oneshot"]
mv --> fut["Future(JoinHandle)"]
st -->|Future| jh[poll JoinHandle]
jh -->|Pending| pend[返回 Pending]
jh -->|拿回服务| back["Service(Some)"]
back --> ask
为什么要有它
PR #283(2019 年,当时是单独的 tower-spawn-ready crate)说明了动机:有些中间件不保证会按时 poll 内层。负载均衡器检查内层是否就绪时是随机挑的,如果某个内层要被 poll 好几次才能就绪,比如它在 poll_ready 里推进一次 TLS 握手,那么均衡器可能永远不会把握手推进到底。SpawnReady 的做法是:没就绪就开一个后台任务专门 poll 它。
问题的根源是 poll_ready 的推进完全靠调用方。调用方 poll 一次、拿到 Pending、然后去忙别的(select! 选了另一个分支,或者均衡器这一轮挑了别的节点),内层就停在那里。内层自己注册的 waker 叫醒的是调用方的任务,调用方醒了也不一定再来问它。
结构:服务或者任务
#[derive(Debug)]
pub struct SpawnReady<S> {
inner: Inner<S>,
}
#[derive(Debug)]
enum Inner<S> {
Service(Option<S>),
Future(tokio::task::JoinHandle<Result<S, BoxError>>),
}
内层服务要么在手里(Service),要么在后台任务里(Future),任务结束时把服务还回来。Option 是为了能在 &mut 下把服务 take 出来移进任务。它没有实现 Clone。
poll_ready:先自己问,问不出来再派人
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), BoxError>> {
loop {
self.inner = match self.inner {
Inner::Service(ref mut svc) => {
if let Poll::Ready(r) = svc.as_mut().expect("illegal state").poll_ready(cx) {
return Poll::Ready(r.map_err(Into::into));
}
let svc = svc.take().expect("illegal state");
let rx =
tokio::spawn(svc.ready_oneshot().map_err(Into::into).in_current_span());
Inner::Future(rx)
}
Inner::Future(ref mut fut) => {
let svc = ready!(Pin::new(fut).poll(cx))??;
Inner::Service(Some(svc))
}
}
}
}
和「把等待放进单独的任务」这句概括相比,有三处细节:
- 第一次是在调用方任务里问的。内层本来就就绪,就直接返回,不开任务。只有返回
Pending才 spawn。 - 移进任务的是服务本身。
ready_oneshot()拿走服务,就绪后把服务作为输出返回(见 ServiceExt 与 Oneshot)。服务在任务里的这段时间,SpawnReady手里什么都没有,poll_ready只能 pollJoinHandle。 - 拿回服务后会再问一次。
Inner::Future分支拿到服务后没有直接返回Ready,而是回到循环顶部,在调用方任务里再poll_ready一次。按 Service trait 的规则,就绪之后再问必须仍是Ready,所以这一次正常情况下立即返回;它保证了SpawnReady返回Ready时,内层刚刚在当前上下文里确认过。
?? 是两层错误:外层是 JoinError(后台任务 panic 或被取消),内层是服务 poll_ready 返回的错误,都转成 BoxError。in_current_span() 让后台任务继承调用方的 tracing span。
因为要 tokio::spawn,内层必须 Send + 'static,而且第一次遇到 Pending 时必须在 tokio 运行时里。
call 只接受服务在手里的情况:
fn call(&mut self, request: Req) -> Self::Future {
match self.inner {
Inner::Service(Some(ref mut svc)) => {
ResponseFuture::new(svc.call(request).map_err(Into::into))
}
_ => unreachable!("poll_ready must be called"),
}
}
服务还在后台任务里时 call,就是没等到 Ready 就调用,unreachable! panic。
实验:只 poll 一次,然后走开
内层模拟一个握手:要被 poll 5 次才就绪,每次 poll 都推进一步,并立刻 wake_by_ref 要求再被 poll。调用方只 poll 一次 poll_ready,然后去睡 10ms:
tower = { version = "=0.5.3", features = ["spawn-ready"] }
tokio = { version = "=1.54.1", features = ["full", "test-util"] }
use std::future::{poll_fn, ready, Ready};
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;
use tower::spawn_ready::SpawnReady;
use tower::{BoxError, Service};
// Needs to be polled 5 times before it is ready, like a handshake
// that makes progress only while someone polls it.
struct Handshake {
polls: Arc<AtomicU32>,
}
impl Service<u32> for Handshake {
type Response = u32;
type Error = BoxError;
type Future = Ready<Result<u32, BoxError>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), BoxError>> {
let n = self.polls.fetch_add(1, Ordering::SeqCst) + 1;
println!(" inner poll_ready #{n} in task {:?}", tokio::task::try_id());
if n < 5 {
cx.waker().wake_by_ref();
return Poll::Pending;
}
Poll::Ready(Ok(()))
}
fn call(&mut self, req: u32) -> Self::Future {
ready(Ok(req))
}
}
// Poll poll_ready exactly once, then go away for a while.
async fn poll_once_then_sleep<S: Service<u32>>(svc: &mut S) -> bool {
let first = poll_fn(|cx| Poll::Ready(svc.poll_ready(cx).is_ready())).await;
tokio::time::sleep(Duration::from_millis(10)).await;
first
}
#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
println!("caller is task {:?}", tokio::task::try_id());
let polls = Arc::new(AtomicU32::new(0));
let mut plain = Handshake { polls: polls.clone() };
println!("plain:");
let first = poll_once_then_sleep(&mut plain).await;
println!(" first poll ready={first}, inner polled {} times", polls.load(Ordering::SeqCst));
let polls = Arc::new(AtomicU32::new(0));
let mut spawned = SpawnReady::new(Handshake { polls: polls.clone() });
println!("SpawnReady:");
let first = poll_once_then_sleep(&mut spawned).await;
println!(" first poll ready={first}, inner polled {} times", polls.load(Ordering::SeqCst));
let again = poll_fn(|cx| Poll::Ready(spawned.poll_ready(cx).is_ready())).await;
println!(" second poll ready={again}, inner polled {} times", polls.load(Ordering::SeqCst));
println!(" call(9) = {}", spawned.call(9).await.unwrap());
}
M2 Max、rustc 1.99.0、--release,连跑 3 次输出相同:
caller is task None
plain:
inner poll_ready #1 in task None
first poll ready=false, inner polled 1 times
SpawnReady:
inner poll_ready #1 in task None
inner poll_ready #2 in task Some(Id(1))
inner poll_ready #3 in task Some(Id(1))
inner poll_ready #4 in task Some(Id(1))
inner poll_ready #5 in task Some(Id(1))
first poll ready=false, inner polled 5 times
inner poll_ready #6 in task None
second poll ready=true, inner polled 6 times
call(9) = 9
task None 是 block_on 直接驱动的 main future,它不是一个 spawn 出来的任务,没有任务 ID。
- plain:调用方 poll 一次就走开,握手停在第 1 步。
- SpawnReady:第 1 次在调用方里问,
Pending,于是服务被移进任务 1。调用方睡觉的时候,任务 1 自己把握手推进到第 5 步,就绪。 - 调用方回来再 poll,
JoinHandle已经完成,拿回服务,在调用方里再问一次(第 6 次),Ready。
被 drop 时取消后台任务
Drop:
impl<S> Drop for SpawnReady<S> {
fn drop(&mut self) {
if let Inner::Future(ref mut task) = self.inner {
task.abort();
}
}
}
服务在后台任务里时 SpawnReady 被丢掉,就 abort 那个任务,服务随任务一起被 drop。abort 怎么生效见 JoinHandle、取消与 JoinError。
Drop 实现是 #581(0.4.7)加的。PR 说明:「when the SpawnReady service is dropped, this task continues to run until it becomes ready (at which point the ready service will be dropped). This can cause resource leaks when the inner service never becomes ready.」也就是说,丢掉 SpawnReady 并不会丢掉它的内层服务:服务在后台任务里,只有等它就绪、任务结束,服务才被释放;它要是永远不就绪,任务和服务就一直留在运行时里。
再早三个月的 #538(0.4.4)把后台任务的回传方式从每次 Pending 新建一个 tokio::sync::oneshot 改成了直接持有 JoinHandle,理由是「This avoid unnecessary allocation」。正是有了 JoinHandle,#581 才能直接 abort。
什么时候用
- 内层的
poll_ready自己在推进一件事(握手、建连接、刷新状态),需要被反复 poll 才能就绪; - 外面的调用方不保证会持续 poll 它,比如它是负载均衡里的一个节点,或者调用方在
select!里同时等别的东西。
代价是每次内层从「就绪」变成「没就绪」都要 spawn 一个任务,服务在两个任务之间来回移动,所以它需要 Send + 'static。内层总是很快就绪、或者调用方本来就会一直 poll 它时,SpawnReady 只是多一层。tower 0.5.3 自己的其他模块都没有用到它。
小结
SpawnReady的状态是Service(Option<S>)或Future(JoinHandle<Result<S, BoxError>>):服务要么在手里,要么在后台任务里。poll_ready先在调用方任务里问内层,Pending才把服务移进tokio::spawn(svc.ready_oneshot());之后只 pollJoinHandle,拿回服务后在调用方任务里再确认一次。- 用途来自 #283:不保证按时 poll 内层的调用方(如负载均衡器)可能永远推进不完内层的就绪过程,后台任务替它推进。
- 服务在后台时
call会unreachable!panic;SpawnReady被 drop 时abort后台任务(#581,0.4.7),永远不就绪的服务不再泄漏。