You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
60 lines
1.6 KiB
Rust
60 lines
1.6 KiB
Rust
use std::io::{ErrorKind, Write};
|
|
use std::sync::mpsc::{sync_channel, Receiver, SyncSender};
|
|
use std::{io, mem};
|
|
|
|
pub struct ChannelSink {
|
|
buffer: Vec<u8>,
|
|
block_size: usize,
|
|
tx: SyncSender<Vec<u8>>,
|
|
}
|
|
|
|
impl ChannelSink {
|
|
/// Creates a new sink with a channel to send the data to
|
|
pub fn new(block_size: usize) -> (Self, Receiver<Vec<u8>>) {
|
|
let (tx, rx) = sync_channel(1);
|
|
(
|
|
Self {
|
|
buffer: Vec::new(),
|
|
block_size,
|
|
tx,
|
|
},
|
|
rx,
|
|
)
|
|
}
|
|
}
|
|
|
|
impl Write for ChannelSink {
|
|
#[tracing::instrument(skip_all, level = "trace")]
|
|
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
|
self.buffer.append(&mut buf.to_vec());
|
|
if self.buffer.len() >= self.block_size {
|
|
tracing::trace!("Block size reached. Sending buffer...");
|
|
self.tx
|
|
.send(mem::take(&mut self.buffer))
|
|
.map_err(|e| io::Error::new(ErrorKind::BrokenPipe, e))?;
|
|
}
|
|
|
|
Ok(buf.len())
|
|
}
|
|
|
|
#[tracing::instrument(skip_all, level = "trace")]
|
|
fn flush(&mut self) -> std::io::Result<()> {
|
|
if !self.buffer.is_empty() {
|
|
self.tx
|
|
.send(mem::take(&mut self.buffer))
|
|
.map_err(|e| io::Error::new(ErrorKind::BrokenPipe, e))?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl Drop for ChannelSink {
|
|
#[tracing::instrument(skip_all, level = "trace")]
|
|
fn drop(&mut self) {
|
|
if let Err(e) = self.flush() {
|
|
tracing::debug!("Error while trying to flush buffer during drop {e}")
|
|
}
|
|
}
|
|
}
|