tonic服务端流式RPC：不使用spawn_blocking执行重计算是否安全？

在实现Tonic的服务端流式RPC时，异步函数需要返回实现Stream trait的类型。如果通过mpsc接收者构造流并调用.map(my_very_expensive_function)，想知道my_very_expensive_function是否会运行在Tokio工作线程上？或者说这个环境是否安全执行阻塞操作？

示例代码：

use proto::MyService;

struct RpcService;

impl MyService for RpcService {
  type MyFnStream = Pin<Box<dyn Stream<Item=Result<Bar, tonic::Status>> + Send>>;

  async fn my_fn(&self, request: tonic::Request<Foo>) -> Result<Response<Self::MyFnStream>, tonic::Status> {
    let (sender, receiver) = tokio::sync::mpsc::channel();
    register_sender(sender);
    Box::new(ReceiverStream(receiver).map(my_very_expensive_function))
  }
}

fn register_sender(sender: tokio::sync::mpsc::Sender<Foo>) {
  ...
}

fn my_very_expensive_function(foo: Foo) -> Bar {
  ... // This should not run in a tokio worker thread
}

答案 #

结论：my_very_expensive_function会运行在Tokio工作线程上，这个环境并不适合执行阻塞操作。

原因 #

Tonic基于Tokio运行时，当客户端从流式RPC拉取数据时，Stream的poll_next方法会在Tokio的工作线程上被调用。而map是同步流操作，它的处理逻辑会直接在调用poll_next的Tokio工作线程上执行。

Tokio工作线程是为异步非阻塞任务设计的，依赖快速释放线程资源来维持调度效率。如果在这里执行长时间阻塞的函数，会占用线程资源，导致其他异步任务无法被及时调度，严重影响服务性能。

解决办法 #

将昂贵的阻塞操作放到Tokio的阻塞线程池中执行，使用tokio::task::spawn_blocking把任务转移到专用线程：

修改后的代码示例：

async fn my_fn(&self, request: tonic::Request<Foo>) -> Result<Response<Self::MyFnStream>, tonic::Status> {
    let (sender, receiver) = tokio::sync::mpsc::channel();
    register_sender(sender);
    
    let stream = ReceiverStream(receiver)
        .then(|foo| async move {
            // 将阻塞操作委托给Tokio阻塞线程池
            let bar = tokio::task::spawn_blocking(move || my_very_expensive_function(foo))
                .await
                .map_err(|e| tonic::Status::internal(format!("Task failed: {}", e)))?;
            Ok(bar)
        });
    
    Ok(Response::new(Box::pin(stream)))
}

这样my_very_expensive_function会运行在Tokio专门为阻塞任务准备的线程池中，不会占用异步工作线程，避免影响服务的整体调度效率。

内容的提问来源于stack exchange，提问作者dspyz