use std::io::{self};
use futures_lite::{AsyncWrite, AsyncWriteExt};
use crate::prelude::*;
use crate::queues::Queue;
pub fn writer_to_bulk_consumer<W, Q>(writer: W, queue: Q) -> WriterToBulkConsumer<W, Q> {
WriterToBulkConsumer::new(writer, queue)
}
#[derive(Debug)]
pub struct WriterToBulkConsumer<W, Q> {
writer: W,
queue: Q,
}
impl<W, Q> WriterToBulkConsumer<W, Q> {
fn new(writer: W, queue: Q) -> Self {
Self { writer, queue }
}
pub fn into_inner(self) -> (W, Q) {
(self.writer, self.queue)
}
}
impl<W, Q> WriterToBulkConsumer<W, Q>
where
W: AsyncWrite + Unpin,
Q: Queue<Item = u8>,
{
async fn flush_internal_buffer(&mut self) -> Result<(), io::Error> {
loop {
let amount = self
.queue
.expose_items(async |items| match self.writer.write_all(items).await {
Ok(()) => (items.len(), Ok(items.len())),
Err(err) => (0, Err(err)),
})
.await?;
if amount == 0 {
break;
}
}
debug_assert!(!self.queue.is_full());
Ok(())
}
}
impl<W, Q> Consumer for WriterToBulkConsumer<W, Q>
where
W: AsyncWrite + Unpin,
Q: Queue<Item = u8>,
{
type Item = u8;
type Final = ();
type Error = io::Error;
async fn consume(&mut self, val: Either<Self::Item, Self::Final>) -> Result<(), Self::Error> {
match val {
Left(item) => {
self.flush_internal_buffer().await?;
self.writer.write_all(&[item]).await
}
Right(()) => {
self.flush_internal_buffer().await?;
self.writer.close().await
}
}
}
async fn flush(&mut self) -> Result<(), Self::Error> {
self.flush_internal_buffer().await?;
self.writer.flush().await
}
}
impl<W, Q> BulkConsumer for WriterToBulkConsumer<W, Q>
where
W: AsyncWrite + Unpin,
Q: Queue<Item = u8>,
{
async fn expose_slots_gracefully<F, R>(&mut self, f: F) -> Result<R, (F, Self::Error)>
where
F: AsyncFnOnce(&mut [Self::Item]) -> (usize, R),
{
if self.queue.is_full() {
if let Err(err) = self.flush_internal_buffer().await {
return Err((f, err));
}
}
Ok(self.queue.expose_slots(async |buffer_slots| {
debug_assert!(!buffer_slots.is_empty(), "A non-full queue must expose at least one item slot when expose_slots is invoked.");
f(buffer_slots).await
}).await)
}
}