minisql/proto/src/reader/utils.rs
2023-12-15 16:21:51 +01:00

24 lines
637 B
Rust

use async_trait::async_trait;
use tokio::io::{AsyncBufRead, AsyncBufReadExt};
#[async_trait]
pub trait AsyncPeek {
async fn peek(&mut self, buf: &mut [u8]) -> tokio::io::Result<usize>;
}
#[async_trait]
impl<T> AsyncPeek for T
where
T: AsyncBufRead + Unpin + Send,
{
async fn peek(&mut self, buf: &mut [u8]) -> tokio::io::Result<usize> {
let filled = self.fill_buf().await?;
if filled.len() >= buf.len() {
buf.copy_from_slice(&filled[..buf.len()]);
Ok(buf.len())
} else {
buf[..filled.len()].copy_from_slice(filled);
Ok(filled.len())
}
}
}