use std::{
pin::Pin,
task::{Context, Poll},
};
use tokio::sync::mpsc;
use tracing::error;
pub(crate) struct ReceiverToAsyncRead {
rx: mpsc::Receiver<String>,
buffer: Vec<u8>,
pos: usize,
}
impl ReceiverToAsyncRead {
pub(crate) fn new(rx: mpsc::Receiver<String>) -> Self {
Self {
rx,
buffer: Vec::new(),
pos: 0,
}
}
}
impl tokio::io::AsyncRead for ReceiverToAsyncRead {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut tokio::io::ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
if self.pos < self.buffer.len() {
let remaining = &self.buffer[self.pos..];
let to_copy = remaining.len().min(buf.remaining());
buf.put_slice(&remaining[..to_copy]);
self.pos += to_copy;
if self.pos >= self.buffer.len() {
self.buffer.clear();
self.pos = 0;
}
return Poll::Ready(Ok(()));
}
match Pin::new(&mut self.rx).poll_recv(cx) {
Poll::Ready(Some(msg)) => {
let bytes = format!("{}\n", msg).into_bytes();
let to_copy = bytes.len().min(buf.remaining());
buf.put_slice(&bytes[..to_copy]);
if to_copy < bytes.len() {
self.buffer = bytes[to_copy..].to_vec();
self.pos = 0;
}
Poll::Ready(Ok(()))
}
Poll::Ready(None) => Poll::Ready(Ok(())),
Poll::Pending => Poll::Pending,
}
}
}
pub(crate) struct SenderToAsyncWrite {
tx: mpsc::Sender<String>,
buffer: Vec<u8>,
}
impl SenderToAsyncWrite {
pub(crate) fn new(tx: mpsc::Sender<String>) -> Self {
Self {
tx,
buffer: Vec::new(),
}
}
}
impl tokio::io::AsyncWrite for SenderToAsyncWrite {
fn poll_write(
mut self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
self.buffer.extend_from_slice(buf);
while let Some(pos) = self.buffer.iter().position(|&b| b == b'\n') {
let line = String::from_utf8_lossy(&self.buffer[..pos]).to_string();
self.buffer.drain(..=pos);
if !line.is_empty() {
if let Err(e) = self.tx.try_send(line.clone()) {
match e {
mpsc::error::TrySendError::Full(_) => {
let truncated: String = line.chars().take(100).collect();
error!(
"Channel full, dropping message (backpressure): {}",
truncated
);
}
mpsc::error::TrySendError::Closed(_) => {
return Poll::Ready(Err(std::io::Error::new(
std::io::ErrorKind::BrokenPipe,
"Channel closed",
)));
}
}
}
}
}
Poll::Ready(Ok(buf.len()))
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Poll::Ready(Ok(()))
}
}