当前位置：首页 > article >正文

Rust实现基于Tokio的限制内存占用的channel

article 2024/9/24 13:22:35

Rust实现基于Tokio的限制内存占用的channel

简介

本文介绍如何基于tokio的channel实现一个限制内存占用的channel。

Tokio提供了多种协程间同步的接口，用于在不同的协程中同步数据。
常用的channel有两种:bounded和unbounded，其中ubbounded的channel可以无限的发送数据,而bounded的channel则有限的发送数据。两种channel都没有对自身的内存占用做出限制。

异步网络编程中常用一个channel连接两个task，其中业务task与业务交互：将要发送的数据发送到channel，而网络task与操作系统交互：从channel中接收数据并写入socket。单有时候带宽有限或者对端接收速率过慢时，而网络task从channel中接收的速度小于业务task向channel中发送的速度时，会造成大量的数据阻塞在channel中，如果不对channel的占用内存做限制，则会造成内存占用过多甚至进程被OOM。

实现

获取数据大小

要想限制channel总的内存占用，必须要直到每个数据的大小。比较常见的作法是所有需要发送到channel的内容都必须实现一个Trait，此Trait中定义了一个get_size方法,用于获取数据的大小。
```
pub trait GetSize {
 /// get total size
 fn get_size(&self) -> usize;
}
```
要发送的内容必须实现GetSize的Trait,并实现get_size方法。注意：get_size方法获取到的大小需包括栈空间和堆空间，例如：
```
 struct MyData {
     data: Vec<u8>,
 }

 impl GetSize for MyData {
     fn get_size(&self) -> usize {
         return std::mem::size_of::<MyData>() + self.data.len();//stack size + heap size
     }
 }
```

创建SizedSender和SizedReceiver

SizedSender和SizedReceiver都可以基于tokio的UnboundedSender和UnboundedReceiver实现。在tokio的基础上，需要共享一个条件变量用于在sender和receiver之间同步当前是否还有可用空间。

   
pub struct SizedSender<T: GetSize> {
    inner: mpsc::UnboundedSender<T>,
    size_semaphore: Arc<(Semaphore, usize)>,
}   


pub struct SizedReceiver<T: GetSize> {
    inner: mpsc::UnboundedReceiver<T>,
    size_semaphore: Arc<(Semaphore, usize)>,
}


/// Limit space usage but not limit the number of messages, bytes_size must bigger than 0.
pub fn sized_channel<T: GetSize>(bytes_size: usize) -> (SizedSender<T>, SizedReceiver<T>) {
    let (tx, rx) = mpsc::unbounded_channel::<T>();
    let semaphore = Arc::new((Semaphore::new(bytes_size), bytes_size));
    (
        SizedSender::new(tx, semaphore.clone()),
        SizedReceiver::new(rx, semaphore),
    )
}

SizedSender实现

发送端发送时需要调用get_size方法获取数据的大小,然后调用Semaphore::available_permits方法获取可用空间,如果可用空间大于数据大小,则发送成功,否则发送失败。

impl<T: GetSize> SizedSender<T> {
 pub fn new(inner: mpsc::UnboundedSender<T>, size_semaphore: Arc<(Semaphore, usize)>) -> Self {
     Self {
         inner,
         size_semaphore,
     }
 }

 fn do_send(
     &self,
     message: T,
     permits: Option<SemaphorePermit<'_>>,
 ) -> Result<(), SendError<T>> {
     match self.inner.send(message) {
         Ok(r) => {
             if let Some(permits) = permits {
                 permits.forget();
             }

             Ok(r)
         }
         Err(e) => {
             log::debug!("send value error!");
             Err(e)
         }
     }
 }
 pub async fn send(&self, message: T) -> Result<(), SendError<T>> {
     let message_size = message.get_size();

     if message_size > self.size_semaphore.1 {
         return Err(SendError(message));
     }
     let size = match u32::try_from(message_size) {
         Ok(size) => size,
         Err(_) => {
             return Err(SendError(message));
         }
     };

     if self.size_semaphore.0.available_permits() < size as usize {
         // The buffer is about to be depleted, sending may be blocked.
     }

     let permits = match self.size_semaphore.0.acquire_many(size).await {
         Ok(perimits) => Some(perimits),
         Err(_) => {
             return Err(SendError(message));
         }
     };

     self.do_send(message, permits)
 }
 }

SizedReceiver的实现

接收端接收时需要调用get_size方法获取数据的大小,然后将相应大小的permits还给信号量即可。

impl<T: GetSize> SizedReceiver<T> {
pub fn new(inner: mpsc::UnboundedReceiver<T>, size_semaphore: Arc<(Semaphore, usize)>) -> Self {
    Self {
        inner,
        size_semaphore,
    }
}

pub async fn recv(&mut self) -> Option<T> {
    self.inner.recv().await.map(|r| {
        let message_size = r.get_size();

        self.size_semaphore.0.add_permits(message_size);

        r
    })
}
}