如何使方法流不可知?

How can I make a method stream agnostic?

我有以下方法:

async fn transfer_all(stream: &mut TcpStream) -> Result<Vec<Vec<u8>>, Box<dyn std::error::Error>> {
  let mut packets: Vec<Vec<u8>> = Vec::new();
  let mut header = true;
  let mut length: usize = 0;

  let mut packet: Vec<u8> = Vec::new();
  loop {
    stream.readable().await?;

    if header {
      length = 5;
      packet.clear();
      packet.shrink_to_fit();
      packet.reserve(length);
    }

    let mut buf: Vec<u8> = vec![0u8; length];
    match stream.try_read(&mut buf) {
      Ok(0) => {
        break;
      }
      Ok(n) => {
        if header {
          length = u32::from_be_bytes(pop(&buf[1..])) as usize - 4;
          header = false;
          packet.append(&mut buf);
          packet.reserve(length);
          continue;
        }
        packet.append(&mut buf);
        packets.push(packet.clone());
        header = true;
      }
      Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
        break;
      }
      Err(e) => {
        return Err(e.into());
      }
    }
  }

  Ok(packets)
}

它适用于 TcpStream,但我还需要使其适用于 UnixStream。由于这是一个相当复杂的状态机,我宁愿没有两个实现。有人建议我使用 async fn transfer_all<S: AsyncRead + Unpin>(stream: &mut S) -> Result<Vec<Vec<u8>>, Box<dyn std::error::Error>> 并将 match stream.try_read(&mut buf) { 替换为 match stream.read(&mut buf).await { 但是当没有更多数据可读时这会阻塞。如何使此方法适用于 TcpStream 和 UnixStream?

由于 UnixStreamTcpStream 都有一个 try_read 方法,你可以为它们创建自己的特征:

trait TryRead {
    // overlapping the name makes it hard to work with
    fn do_try_read(&self, buf: &mut [u8]) -> Result<usize>;
}

impl TryRead for TcpStream {
    fn do_try_read(&self, buf: &mut [u8]) -> Result<usize> {
        self.try_read(buf)
    }
}

impl TryRead for UnixStream {
    fn do_try_read(&self, buf: &mut [u8]) -> Result<usize> {
        self.try_read(buf)
    }
}

然后,你可以取一个S: AsyncRead + TryRead + Unpin,然后用do_try_read替换try_read

或者,使用减少重复的宏:

trait TryRead {
    // overlapping the name makes it hard to work with
    fn do_try_read(&self, buf: &mut [u8]) -> Result<usize>;
}

macro_rules! make_try_read {
    ($typ: ty) => {
        impl TryRead for $typ {
            fn do_try_read(&self, buf: &mut [u8]) -> Result<usize> {
                self.try_read(buf)
            }
        }
    }
}

make_try_read!(TcpStream);
make_try_read!(UnixStream);