use std::io;
use std::path::Path;
use std::pin::Pin;
use std::task::{Context, Poll};
use futures_core::Stream;
use minarrow::structs::shared_buffer::SharedBuffer;
use tokio::io::{AsyncRead, ReadBuf};
use tokio::net::unix::OwnedReadHalf;
use crate::enums::BufferChunkSize;
use crate::models::streams::stream_arena::StreamArena;
use crate::models::transports::uds::UdsTransport;
pub struct UdsByteStream {
reader: OwnedReadHalf,
eof: bool,
chunk_size: usize,
arena: StreamArena,
}
impl UdsByteStream {
pub async fn connect(path: impl AsRef<Path>) -> io::Result<Self> {
let (read_half, _write_half) = UdsTransport::connect(path).await?;
Ok(Self::from_read_half(read_half, BufferChunkSize::Http))
}
pub fn from_read_half(read_half: OwnedReadHalf, size: BufferChunkSize) -> Self {
Self {
reader: read_half,
eof: false,
chunk_size: size.chunk_size(),
arena: StreamArena::new(),
}
}
}
impl Stream for UdsByteStream {
type Item = Result<SharedBuffer, io::Error>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let me = self.get_mut();
if me.eof {
return Poll::Ready(None);
}
if me.arena.remaining() < me.chunk_size {
me.arena.recycle_or_reset();
}
let chunk_start = me.arena.write_pos();
let n = {
let spare = me.arena.spare_uninit();
let read_len = spare.len().min(me.chunk_size);
let mut read_buf = ReadBuf::uninit(&mut spare[..read_len]);
match Pin::new(&mut me.reader).poll_read(cx, &mut read_buf) {
Poll::Ready(Ok(())) => read_buf.filled().len(),
Poll::Ready(Err(e)) => {
me.eof = true;
return Poll::Ready(Some(Err(e)));
}
Poll::Pending => return Poll::Pending,
}
};
if n == 0 {
me.eof = true;
return Poll::Ready(None);
}
unsafe { me.arena.advance(n) };
let shared = me.arena.window(chunk_start, n);
me.arena.align();
Poll::Ready(Some(Ok(shared)))
}
}
impl AsyncRead for UdsByteStream {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let me = self.get_mut();
Pin::new(&mut me.reader).poll_read(cx, buf)
}
}