diff --git a/attic/src/nix_store/bindings/mod.rs b/attic/src/nix_store/bindings/mod.rs index 7487ee01..ef1df694 100644 --- a/attic/src/nix_store/bindings/mod.rs +++ b/attic/src/nix_store/bindings/mod.rs @@ -43,11 +43,14 @@ pub unsafe fn open_nix_store() -> AtticResult { // (tokio, crossbeam, flume) mod mpsc { // Tokio - pub use tokio::sync::mpsc::{ - UnboundedReceiver, UnboundedSender, error::SendError, unbounded_channel, - }; + pub use tokio::sync::mpsc::{Receiver, Sender, channel, error::SendError}; } +/// Maximum number of NAR chunks buffered between the C++ serializer and the uploader. +const NAR_CHANNEL_CAPACITY: usize = 64; + +const NAR_CHUNK_SIZE: usize = 64 * 1024; + /// Async write request. #[derive(Debug)] enum AsyncWriteMessage { @@ -59,44 +62,76 @@ enum AsyncWriteMessage { /// Async write request sender. #[derive(Clone)] pub struct AsyncWriteSender { - sender: mpsc::UnboundedSender, + sender: mpsc::Sender, + buf: Vec, } impl AsyncWriteSender { fn send(&mut self, data: &[u8]) -> Result<(), mpsc::SendError> { - let message = AsyncWriteMessage::Data(Vec::from(data)); - self.sender.send(message) + let mut rest = data; + + while !rest.is_empty() { + if self.buf.capacity() == 0 { + self.buf.reserve_exact(NAR_CHUNK_SIZE); + } + + let take = (NAR_CHUNK_SIZE - self.buf.len()).min(rest.len()); + + self.buf.extend_from_slice(&rest[..take]); + rest = &rest[take..]; + + if self.buf.len() == NAR_CHUNK_SIZE { + self.flush()?; + } + } + + Ok(()) } fn eof(&mut self) -> Result<(), mpsc::SendError> { - let message = AsyncWriteMessage::Eof; - self.sender.send(message) + self.flush()?; + self.sender.blocking_send(AsyncWriteMessage::Eof) } pub(crate) fn rust_error( &mut self, error: impl std::error::Error, ) -> Result<(), impl std::error::Error> { + // No need to preserve buffered data when the consumer is about to abort. let message = AsyncWriteMessage::Error(error.to_string()); - self.sender.send(message) + self.sender.blocking_send(message) + } + + fn flush(&mut self) -> Result<(), mpsc::SendError> { + if self.buf.is_empty() { + return Ok(()); + } + + // This runs on the spawn_blocking thread driving the C++ serializer, + // so blocking here does not stall a Tokio runtime worker. + let chunk = std::mem::take(&mut self.buf); + self.sender.blocking_send(AsyncWriteMessage::Data(chunk)) } } /// A wrapper of the `AsyncWrite` trait for the synchronous Nix C++ land. pub struct AsyncWriteAdapter { - receiver: mpsc::UnboundedReceiver, + receiver: mpsc::Receiver, eof: bool, } impl AsyncWriteAdapter { pub fn new() -> (Self, Box) { - let (sender, receiver) = mpsc::unbounded_channel(); + let (sender, receiver) = mpsc::channel(NAR_CHANNEL_CAPACITY); let r = Self { receiver, eof: false, }; - let sender = Box::new(AsyncWriteSender { sender }); + let sender = Box::new(AsyncWriteSender { + sender, + buf: Vec::new(), + }); (r, sender) }