如何在tonic中从客户端施加背压，避免过快消费流式RPC？

问题描述 #

我实现了一个服务器端流式RPC服务，原本设计为根据客户端.next()的调用速度丢弃中间更新（仅返回最新结果），但实际服务器会持续循环调用.next()，完全不受客户端消费速度的影响。请问如何从客户端施加背压，让.next()的调用速率受限于客户端的消费能力？

我尝试在客户端设置buffer_size，但没有任何效果，也找不到服务器端的对应设置项。

场景复现：服务器每秒递增计数器并向订阅者推送数值，客户端每2秒才消费一次消息。运行后客户端会无限落后，队列似乎完全无界。即使客户端消费速度快，通过toxiproxy限流连接时也会出现同样情况，说明这并非客户端缓存无界导致的。

示例代码 #

定义proto文件 #

// streaming.proto

syntax = "proto3";

package streaming;

service Counter {
  rpc StreamCounter(Empty) returns (stream CounterResponse);
}

message Empty {}

message CounterResponse {
  int32 value = 1;
}

服务器端代码（原实现） #

// server.rs

use futures_util::stream;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::watch;
use tokio::time::{sleep, Duration};
use tokio_stream::Stream;
use tonic::async_trait;
use tonic::{transport::Server, Request, Response, Status};

pub mod streaming {
    tonic::include_proto!("streaming");
}

use streaming::counter_server::{Counter, CounterServer};
use streaming::{CounterResponse, Empty};

pub struct MyCounter {
    tx: Arc<watch::Sender<i32>>,
}

#[async_trait]
impl Counter for MyCounter {
    type StreamCounterStream = Pin<Box<dyn Stream<Item = Result<CounterResponse, Status>> + Send>>;

    async fn stream_counter(
        &self,
        _request: Request<Empty>,
    ) -> Result<Response<Self::StreamCounterStream>, Status> {
        let rx = self.tx.subscribe();

        let output = stream::unfold(rx, |mut rx| async move {
            let result = rx.changed().await;
            match result {
                Ok(()) => {
                    let value = *rx.borrow();
                    println!("Sending: {}", value);
                    Some((Ok(CounterResponse { value }), rx))
                }
                Err(_) => None,
            }
        });

        Ok(Response::new(Box::pin(output)))
    }
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let (tx, _rx) = watch::channel(0);
    let tx = Arc::new(tx);
    let my_counter = MyCounter { tx: tx.clone() };

    tokio::spawn(async move {
        let mut counter = 0;
        loop {
            sleep(Duration::from_secs(1)).await;
            counter += 1;
            let _ = tx.send(counter);
        }
    });

    let addr = "[::1]:50051".parse()?;
    Server::builder()
        .add_service(CounterServer::new(my_counter))
        .serve(addr)
        .await?;

    Ok(())
}

客户端代码（原实现） #

// client.rs

use streaming::counter_client::CounterClient;
use streaming::Empty;
use tokio_stream::StreamExt;
use tonic::transport::Endpoint;
use tonic::Request;

pub mod streaming {
    tonic::include_proto!("streaming");
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut client =
        CounterClient::connect(Endpoint::from_static("http://[::1]:50051").buffer_size(Some(2)))
            .await?;
    let request = Request::new(Empty {});

    let mut stream = client.stream_counter(request).await?.into_inner();

    while let Some(response) = stream.next().await {
        match response {
            Ok(counter_response) => {
                println!("Received: {}", counter_response.value);
                tokio::time::sleep(tokio::time::Duration::from_secs(2)).await; // 模拟慢处理
            }
            Err(e) => {
                eprintln!("Error: {}", e);
                break;
            }
        }
    }

    Ok(())
}

解决方案 #

问题的核心在于原服务器端使用的watch channel不支持背压机制——它只会持续推送最新值，完全不关心客户端是否已经消费了之前的消息。要实现客户端驱动的背压，需要替换为支持背压的通信方式，并在服务器端处理消息发送的逻辑：

1. 使用Tokio的mpsc channel替代watch：mpsc支持背压，当接收端（客户端侧的stream）消费速度慢时，发送端会被阻塞，或者可以通过try_send丢弃旧消息。
2. 服务器端stream从mpsc receiver生成：这样stream的.next()调用会等待客户端准备好接收，自然实现速率限制。
3. 若需要丢弃中间更新，在发送新值前检查channel是否有未被消费的消息，有则丢弃旧值发送新值。

修改后的服务器端代码 #

// server.rs
use futures_util::StreamExt;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio::time::{sleep, Duration};
use tokio_stream::Stream;
use tonic::async_trait;
use tonic::{transport::Server, Request, Response, Status};

pub mod streaming {
    tonic::include_proto!("streaming");
}

use streaming::counter_server::{Counter, CounterServer};
use streaming::{CounterResponse, Empty};

pub struct MyCounter {
    tx: Arc<mpsc::Sender<i32>>,
}

#[async_trait]
impl Counter for MyCounter {
    type StreamCounterStream = Pin<Box<dyn Stream<Item = Result<CounterResponse, Status>> + Send>>;

    async fn stream_counter(
        &self,
        _request: Request<Empty>,
    ) -> Result<Response<Self::StreamCounterStream>, Status> {
        // 为每个客户端创建独立的receiver
        let (tx_client, rx_client) = mpsc::channel(1); // buffer size设为1，保留最新值
        let tx_clone = self.tx.clone();

        // 启动任务将全局计数器转发给客户端receiver
        tokio::spawn(async move {
            let mut rx_global = tx_clone.subscribe();
            while let Some(value) = rx_global.recv().await {
                // 如果客户端还没消费上一个值，就丢弃旧值发送新值
                let _ = tx_client.try_send(value);
            }
        });

        // 将client receiver转换为tonic需要的stream
        let output = rx_client.map(|value| Ok(CounterResponse { value }));

        Ok(Response::new(Box::pin(output)))
    }
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 创建全局mpsc channel，buffer size设为1
    let (tx, _rx) = mpsc::channel(1);
    let tx = Arc::new(tx);
    let my_counter = MyCounter { tx: tx.clone() };

    tokio::spawn(async move {
        let mut counter = 0;
        loop {
            sleep(Duration::from_secs(1)).await;
            counter += 1;
            // 若全局channel满，丢弃旧值发送新值
            let _ = tx.try_send(counter);
        }
    });

    let addr = "[::1]:50051".parse()?;
    Server::builder()
        .add_service(CounterServer::new(my_counter))
        .serve(addr)
        .await?;

    Ok(())
}

关键改动说明 #

• 全局计数器使用mpsc channel，try_send确保只保留最新值，旧值会被丢弃。
• 每个客户端连接时创建独立的mpsc channel，全局消息转发任务会将最新值推送给客户端，若客户端未消费则丢弃旧值。
• 客户端侧的stream从mpsc receiver生成，当客户端消费慢时，receiver的recv()会等待，进而服务器的转发任务会被try_send阻塞（或丢弃旧值），实现背压。

客户端无需修改 #

原客户端代码可以直接使用，当客户端每2秒消费一次时，服务器只会每2秒推送一次最新的计数器值，不会出现无限落后的情况。

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