diff --git c/attic/src/nix_store/bindings/mod.rs i/attic/src/nix_store/bindings/mod.rs index a16b52a..37faf08 100644 --- c/attic/src/nix_store/bindings/mod.rs +++ i/attic/src/nix_store/bindings/mod.rs @@ -10,6 +10,8 @@ use tokio::io::{AsyncWrite, AsyncWriteExt}; use crate::{AtticError, AtticResult}; +use super::stream_config::NarStreamConfig; + // The C++ implementation takes care of concurrency #[repr(transparent)] pub struct FfiNixStore(UnsafeCell>); @@ -42,10 +44,12 @@ pub unsafe fn open_nix_store() -> AtticResult { mod mpsc { // Tokio pub use tokio::sync::mpsc::{ - error::SendError, unbounded_channel, UnboundedReceiver, UnboundedSender, + error::SendError, channel, Receiver, Sender, }; } + + /// Async write request. #[derive(Debug)] enum AsyncWriteMessage { @@ -57,18 +61,20 @@ enum AsyncWriteMessage { /// Async write request sender. #[derive(Clone)] pub struct AsyncWriteSender { - sender: mpsc::UnboundedSender, + sender: mpsc::Sender, } impl AsyncWriteSender { fn send(&mut self, data: &[u8]) -> Result<(), mpsc::SendError> { let message = AsyncWriteMessage::Data(Vec::from(data)); - self.sender.send(message) + // Use blocking_send since we're called from C++ in a blocking context + // This provides backpressure when the channel is full + self.sender.blocking_send(message).map_err(|e| mpsc::SendError(e.0)) } fn eof(&mut self) -> Result<(), mpsc::SendError> { let message = AsyncWriteMessage::Eof; - self.sender.send(message) + self.sender.blocking_send(message).map_err(|e| mpsc::SendError(e.0)) } pub(crate) fn rust_error( @@ -76,19 +82,25 @@ impl AsyncWriteSender { error: impl std::error::Error, ) -> Result<(), impl std::error::Error> { let message = AsyncWriteMessage::Error(error.to_string()); - self.sender.send(message) + // Try to send error without blocking, it's OK if channel is closed + self.sender.try_send(message) } } /// 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(); + Self::new_with_config(NarStreamConfig::default()) + } + + pub fn new_with_config(config: NarStreamConfig) -> (Self, Box) { + // Use bounded channel to provide backpressure and prevent OOM + let (sender, receiver) = mpsc::channel(config.channel_capacity); let r = Self { receiver, diff --git c/attic/src/nix_store/mod.rs i/attic/src/nix_store/mod.rs index 4e08a67..04ef0f6 100644 --- c/attic/src/nix_store/mod.rs +++ i/attic/src/nix_store/mod.rs @@ -46,6 +46,9 @@ mod bindings; #[cfg(feature = "nix_store")] mod nix_store; +#[cfg(feature = "nix_store")] +mod stream_config; + use std::ffi::OsStr; #[cfg(target_family = "unix")] use std::os::unix::ffi::OsStrExt; @@ -61,6 +64,9 @@ use crate::hash::Hash; #[cfg(feature = "nix_store")] pub use nix_store::NixStore; +#[cfg(feature = "nix_store")] +pub use stream_config::NarStreamConfig; + #[cfg(test)] pub mod tests; diff --git c/attic/src/nix_store/nix_store.rs i/attic/src/nix_store/nix_store.rs index 3b754f9..91eb489 100644 --- c/attic/src/nix_store/nix_store.rs +++ i/attic/src/nix_store/nix_store.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use tokio::task::spawn_blocking; use super::bindings::{open_nix_store, AsyncWriteAdapter, FfiNixStore}; +use super::stream_config::NarStreamConfig; use super::{to_base_name, StorePath, ValidPathInfo}; use crate::error::AtticResult; use crate::hash::Hash; @@ -19,11 +20,18 @@ pub struct NixStore { /// Path to the Nix store itself. store_dir: PathBuf, + + /// Configuration for NAR streaming. + stream_config: NarStreamConfig, } #[cfg(feature = "nix_store")] impl NixStore { pub fn connect() -> AtticResult { + Self::connect_with_config(NarStreamConfig::from_env()) + } + + pub fn connect_with_config(stream_config: NarStreamConfig) -> AtticResult { #[allow(unsafe_code)] let inner = unsafe { open_nix_store()? }; let store_dir = PathBuf::from(inner.store().store_dir()); @@ -31,6 +39,7 @@ impl NixStore { Ok(Self { inner: Arc::new(inner), store_dir, + stream_config, }) } @@ -84,12 +93,26 @@ impl NixStore { self.store_dir.join(&store_path.base_name) } - /// Creates a NAR archive from a path. + /// Creates a NAR archive from a path with default configuration. /// /// This is akin to `nix-store --dump`. pub fn nar_from_path(&self, store_path: StorePath) -> AsyncWriteAdapter { + self.nar_from_path_with_config(store_path, self.stream_config) + } + + /// Creates a NAR archive from a path with custom streaming configuration. + /// + /// This is akin to `nix-store --dump`. + /// + /// The `stream_config` parameter allows you to control the memory buffer size + /// for the NAR stream, which is important for large files to prevent OOM. + pub fn nar_from_path_with_config( + &self, + store_path: StorePath, + stream_config: NarStreamConfig, + ) -> AsyncWriteAdapter { let inner = self.inner.clone(); - let (adapter, mut sender) = AsyncWriteAdapter::new(); + let (adapter, mut sender) = AsyncWriteAdapter::new_with_config(stream_config); let base_name = Vec::from(store_path.as_base_name_bytes()); spawn_blocking(move || { diff --git c/attic/src/nix_store/stream_config.rs i/attic/src/nix_store/stream_config.rs new file mode 100644 index 0000000..bdc6177 --- /dev/null +++ i/attic/src/nix_store/stream_config.rs @@ -0,0 +1,116 @@ +//! Configuration for NAR streaming. +//! +//! This module provides configuration options for controlling memory usage +//! during NAR stream uploads. + +use std::env; + +/// Default capacity for the NAR streaming channel. +/// +/// At ~256KB per chunk (typical), this allows: +/// - 64 chunks = 16 MB buffer (conservative) +/// - 128 chunks = 32 MB buffer (balanced, default) +/// - 256 chunks = 64 MB buffer (high throughput) +pub const DEFAULT_NAR_STREAM_CHANNEL_CAPACITY: usize = 128; + +/// Minimum allowed channel capacity to prevent starvation. +pub const MIN_NAR_STREAM_CHANNEL_CAPACITY: usize = 8; + +/// Maximum allowed channel capacity to prevent excessive memory usage. +pub const MAX_NAR_STREAM_CHANNEL_CAPACITY: usize = 1024; + +/// Configuration for NAR streaming behavior. +#[derive(Debug, Clone, Copy)] +pub struct NarStreamConfig { + /// Number of chunks to buffer in the NAR streaming channel. + /// + /// This provides backpressure to prevent the C++ NAR generator from + /// getting too far ahead of network upload, which could cause OOM + /// when uploading large files. + pub channel_capacity: usize, +} + +impl Default for NarStreamConfig { + fn default() -> Self { + Self { + channel_capacity: DEFAULT_NAR_STREAM_CHANNEL_CAPACITY, + } + } +} + +impl NarStreamConfig { + /// Creates a new configuration with the default values. + pub fn new() -> Self { + Self::default() + } + + /// Sets the channel capacity. + /// + /// The value will be clamped to the valid range. + pub fn with_channel_capacity(mut self, capacity: usize) -> Self { + self.channel_capacity = capacity.clamp( + MIN_NAR_STREAM_CHANNEL_CAPACITY, + MAX_NAR_STREAM_CHANNEL_CAPACITY, + ); + self + } + + /// Loads configuration from environment variables. + /// + /// Supported environment variables: + /// - `ATTIC_NAR_STREAM_CHANNEL_CAPACITY`: Channel capacity (default: 128) + pub fn from_env() -> Self { + let mut config = Self::default(); + + if let Ok(capacity_str) = env::var("ATTIC_NAR_STREAM_CHANNEL_CAPACITY") { + if let Ok(capacity) = capacity_str.parse::() { + config = config.with_channel_capacity(capacity); + } + } + + config + } + + /// Estimates the maximum memory usage for this configuration. + /// + /// This assumes an average chunk size. Actual usage may vary. + pub fn estimated_max_memory_bytes(&self, avg_chunk_size: usize) -> usize { + self.channel_capacity * avg_chunk_size + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_default_config() { + let config = NarStreamConfig::default(); + assert_eq!(config.channel_capacity, DEFAULT_NAR_STREAM_CHANNEL_CAPACITY); + } + + #[test] + fn test_with_channel_capacity() { + let config = NarStreamConfig::new().with_channel_capacity(256); + assert_eq!(config.channel_capacity, 256); + } + + #[test] + fn test_channel_capacity_clamping() { + // Test minimum clamping + let config = NarStreamConfig::new().with_channel_capacity(1); + assert_eq!(config.channel_capacity, MIN_NAR_STREAM_CHANNEL_CAPACITY); + + // Test maximum clamping + let config = NarStreamConfig::new().with_channel_capacity(10000); + assert_eq!(config.channel_capacity, MAX_NAR_STREAM_CHANNEL_CAPACITY); + } + + #[test] + fn test_estimated_memory() { + let config = NarStreamConfig::new().with_channel_capacity(128); + let chunk_size = 256 * 1024; // 256 KB + let estimated = config.estimated_max_memory_bytes(chunk_size); + assert_eq!(estimated, 128 * 256 * 1024); // 32 MB + } +} diff --git c/attic/src/nix_store/stream_test.rs i/attic/src/nix_store/stream_test.rs new file mode 100644 index 0000000..e9c18f2 --- /dev/null +++ i/attic/src/nix_store/stream_test.rs @@ -0,0 +1,304 @@ +//! Tests for NAR streaming with bounded channels. + +#[cfg(test)] +mod tests { + use super::super::bindings::AsyncWriteAdapter; + use super::super::stream_config::NarStreamConfig; + use futures::StreamExt; + use std::time::Duration; + use tokio::time::timeout; + + /// Test that bounded channel provides backpressure. + #[tokio::test] + async fn test_bounded_channel_backpressure() { + let config = NarStreamConfig::new().with_channel_capacity(2); + let (mut adapter, mut sender) = AsyncWriteAdapter::new_with_config(config); + + // Send data in a separate task + let send_handle = tokio::spawn(async move { + // These should succeed immediately (capacity = 2) + sender.send(&[1u8; 100]).unwrap(); + sender.send(&[2u8; 100]).unwrap(); + + // This should block until receiver consumes some data + // We spawn this in a blocking context since send uses blocking_send + let result = tokio::task::spawn_blocking(move || { + sender.send(&[3u8; 100]) + }).await; + + result.unwrap() + }); + + // Give sender time to fill the channel + tokio::time::sleep(Duration::from_millis(100)).await; + + // Now start consuming + let first = adapter.next().await.unwrap().unwrap(); + assert_eq!(first.len(), 100); + assert_eq!(first[0], 1); + + let second = adapter.next().await.unwrap().unwrap(); + assert_eq!(second.len(), 100); + assert_eq!(second[0], 2); + + // The third send should now complete + let send_result = timeout(Duration::from_secs(1), send_handle).await; + assert!(send_result.is_ok()); + + let third = adapter.next().await.unwrap().unwrap(); + assert_eq!(third.len(), 100); + assert_eq!(third[0], 3); + } + + /// Test that channel capacity is enforced. + #[tokio::test] + async fn test_channel_capacity_enforcement() { + let config = NarStreamConfig::new().with_channel_capacity(4); + let (mut adapter, mut sender) = AsyncWriteAdapter::new_with_config(config); + + // Fill the channel completely + tokio::task::spawn_blocking(move || { + for i in 0..4 { + sender.send(&[i as u8; 50]).unwrap(); + } + sender.eof().unwrap(); + }); + + // Give it time to fill + tokio::time::sleep(Duration::from_millis(50)).await; + + // Consume all items + let mut count = 0; + while let Some(Ok(chunk)) = adapter.next().await { + assert_eq!(chunk.len(), 50); + assert_eq!(chunk[0], count); + count += 1; + } + + assert_eq!(count, 4); + } + + /// Test that small channel capacity works correctly. + #[tokio::test] + async fn test_minimum_channel_capacity() { + let config = NarStreamConfig::new().with_channel_capacity(8); // minimum + let (mut adapter, mut sender) = AsyncWriteAdapter::new_with_config(config); + + let send_handle = tokio::task::spawn_blocking(move || { + for i in 0..16 { + sender.send(&[i as u8; 10]).unwrap(); + } + sender.eof().unwrap(); + }); + + let mut received = Vec::new(); + while let Some(Ok(chunk)) = adapter.next().await { + received.push(chunk[0]); + } + + send_handle.await.unwrap(); + assert_eq!(received.len(), 16); + assert_eq!(received, (0..16).collect::>()); + } + + /// Test that large channel capacity works correctly. + #[tokio::test] + async fn test_large_channel_capacity() { + let config = NarStreamConfig::new().with_channel_capacity(256); + let (mut adapter, mut sender) = AsyncWriteAdapter::new_with_config(config); + + let send_handle = tokio::task::spawn_blocking(move || { + for i in 0..100 { + sender.send(&[i as u8; 1024]).unwrap(); + } + sender.eof().unwrap(); + }); + + let mut count = 0; + while let Some(Ok(chunk)) = adapter.next().await { + assert_eq!(chunk.len(), 1024); + assert_eq!(chunk[0], count); + count += 1; + } + + send_handle.await.unwrap(); + assert_eq!(count, 100); + } + + /// Test error handling in bounded channel. + #[tokio::test] + async fn test_error_handling() { + let config = NarStreamConfig::new().with_channel_capacity(2); + let (mut adapter, mut sender) = AsyncWriteAdapter::new_with_config(config); + + tokio::task::spawn_blocking(move || { + sender.send(&[1u8; 100]).unwrap(); + sender.rust_error(std::io::Error::new( + std::io::ErrorKind::Other, + "Test error" + )).ok(); // try_send, might fail if channel closed + }); + + // First message should be data + let first = adapter.next().await.unwrap().unwrap(); + assert_eq!(first[0], 1); + + // Second message should be error + let second = adapter.next().await.unwrap(); + assert!(second.is_err()); + } + + /// Test that EOF is properly handled. + #[tokio::test] + async fn test_eof_handling() { + let config = NarStreamConfig::new().with_channel_capacity(4); + let (mut adapter, mut sender) = AsyncWriteAdapter::new_with_config(config); + + tokio::task::spawn_blocking(move || { + sender.send(&[1u8; 100]).unwrap(); + sender.send(&[2u8; 100]).unwrap(); + sender.eof().unwrap(); + }); + + let first = adapter.next().await.unwrap().unwrap(); + assert_eq!(first[0], 1); + + let second = adapter.next().await.unwrap().unwrap(); + assert_eq!(second[0], 2); + + // Should be EOF now + assert!(adapter.next().await.is_none()); + } + + /// Test concurrent producers (should not happen in practice, but test safety). + #[tokio::test] + async fn test_concurrent_safety() { + let config = NarStreamConfig::new().with_channel_capacity(16); + let (mut adapter, sender) = AsyncWriteAdapter::new_with_config(config); + + // Spawn multiple producers + let mut handles = Vec::new(); + for i in 0..4 { + let mut sender_clone = sender.clone(); + let handle = tokio::task::spawn_blocking(move || { + for j in 0..10 { + let val = (i * 10 + j) as u8; + sender_clone.send(&[val; 50]).unwrap(); + } + }); + handles.push(handle); + } + + // Wait for all producers + for handle in handles { + handle.await.unwrap(); + } + + // Signal EOF from one sender + let mut sender = *sender; + tokio::task::spawn_blocking(move || { + sender.eof().unwrap(); + }); + + // Consume all data + let mut count = 0; + while let Some(Ok(_chunk)) = adapter.next().await { + count += 1; + } + + assert_eq!(count, 40); // 4 producers * 10 chunks each + } + + /// Benchmark-style test: ensure bounded channel doesn't significantly + /// slow down high-throughput scenarios. + #[tokio::test] + async fn test_throughput_performance() { + let config = NarStreamConfig::new().with_channel_capacity(128); + let (mut adapter, mut sender) = AsyncWriteAdapter::new_with_config(config); + + let start = std::time::Instant::now(); + + // Simulate high-throughput upload + let send_handle = tokio::task::spawn_blocking(move || { + for i in 0..1000 { + // 256KB chunks + let chunk = vec![i as u8; 256 * 1024]; + sender.send(&chunk).unwrap(); + } + sender.eof().unwrap(); + }); + + // Consume as fast as possible + let mut total_bytes = 0u64; + while let Some(Ok(chunk)) = adapter.next().await { + total_bytes += chunk.len() as u64; + } + + send_handle.await.unwrap(); + let elapsed = start.elapsed(); + + // Should process 256 MB in reasonable time (< 5 seconds even on slow systems) + assert_eq!(total_bytes, 1000 * 256 * 1024); + assert!(elapsed.as_secs() < 5, "Throughput test too slow: {:?}", elapsed); + + let throughput_mbps = (total_bytes as f64 / 1024.0 / 1024.0) / elapsed.as_secs_f64(); + println!("Throughput: {:.2} MB/s", throughput_mbps); + } + + /// Test memory-bounded behavior: ensure we don't allocate unbounded memory. + #[tokio::test] + async fn test_memory_bounded() { + let config = NarStreamConfig::new().with_channel_capacity(10); + let (mut adapter, mut sender) = AsyncWriteAdapter::new_with_config(config); + + // Try to send more than channel capacity + let send_handle = tokio::task::spawn_blocking(move || { + // This would OOM with unbounded channel on large iterations + // With bounded channel, it will block and wait + for i in 0..100 { + let chunk = vec![i as u8; 1024 * 1024]; // 1 MB chunks + sender.send(&chunk).unwrap(); + } + sender.eof().unwrap(); + }); + + // Slow consumer + let mut count = 0; + while let Some(Ok(_chunk)) = adapter.next().await { + count += 1; + // Simulate slow network + tokio::time::sleep(Duration::from_millis(1)).await; + } + + send_handle.await.unwrap(); + assert_eq!(count, 100); + } + + /// Test configuration edge cases. + #[test] + fn test_config_clamping() { + // Below minimum + let config = NarStreamConfig::new().with_channel_capacity(1); + assert_eq!(config.channel_capacity, 8); // clamped to minimum + + // Above maximum + let config = NarStreamConfig::new().with_channel_capacity(10000); + assert_eq!(config.channel_capacity, 1024); // clamped to maximum + + // Within range + let config = NarStreamConfig::new().with_channel_capacity(100); + assert_eq!(config.channel_capacity, 100); + } + + /// Test estimated memory calculation. + #[test] + fn test_memory_estimation() { + let config = NarStreamConfig::new().with_channel_capacity(128); + let chunk_size = 256 * 1024; // 256 KB + let estimated = config.estimated_max_memory_bytes(chunk_size); + + // 128 chunks * 256 KB = 32 MB + assert_eq!(estimated, 128 * 256 * 1024); + assert_eq!(estimated, 33_554_432); // 32 MB + } +}