Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 44 additions & 12 deletions nativelink-store/tests/filesystem_store_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -363,6 +363,35 @@ const HASH1: &str = "0123456789abcdef000000000000000000010000000000000123456789a
const HASH2: &str = "0123456789abcdef000000000000000000020000000000000123456789abcdef";
const VALUE1: &str = "0123456789";
const VALUE2: &str = "9876543210";

// Longer, in `read_buffer_size: 1` chunks, than `make_buf_channel_pair`'s
// in-flight capacity, so a reader that has taken only the first byte keeps the
// store's `get` stalled with the file still open. Tests that need an open
// stream on a file use these instead of `VALUE1`/`VALUE2`.
const LONG_VALUE1: &str = concat!(
"0123456789",
"0123456789",
"0123456789",
"0123456789",
"0123456789",
"0123456789",
"0123456789",
"0123456789",
"0123456789",
"0123456789",
);
const LONG_VALUE2: &str = concat!(
"9876543210",
"9876543210",
"9876543210",
"9876543210",
"9876543210",
"9876543210",
"9876543210",
"9876543210",
"9876543210",
"9876543210",
);
const STRING_NAME: &str = "String_Filename";

#[nativelink_test]
Expand Down Expand Up @@ -523,7 +552,7 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error>
}
}

let digest1 = DigestInfo::try_new(HASH1, VALUE1.len())?;
let digest1 = DigestInfo::try_new(HASH1, LONG_VALUE1.len())?;
let content_path = make_temp_path("content_path");
let temp_path = make_temp_path("temp_path");

Expand All @@ -543,7 +572,7 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error>
);

// Insert data into store.
store.update_oneshot(digest1, VALUE1.into()).await?;
store.update_oneshot(digest1, LONG_VALUE1.into()).await?;

let (writer, mut reader) = make_buf_channel_pair();
let store_clone = store.clone();
Expand All @@ -561,13 +590,13 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error>
.err_tip(|| "Error reading first byte")?;
assert_eq!(
first_byte[0],
VALUE1.as_bytes()[0],
LONG_VALUE1.as_bytes()[0],
"Expected first byte to match"
);
}

// Replace content.
store.update_oneshot(digest1, VALUE2.into()).await?;
store.update_oneshot(digest1, LONG_VALUE2.into()).await?;

// Ensure we let any background tasks finish.
tokio::task::yield_now().await;
Expand All @@ -586,7 +615,7 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error>
let data = read_file_contents(path.as_os_str()).await?;
assert_eq!(
&data[..],
VALUE1.as_bytes(),
LONG_VALUE1.as_bytes(),
"Expected file content to match"
);
}
Expand All @@ -603,7 +632,7 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error>

assert_eq!(
&remaining_file_data,
&VALUE1.as_bytes()[1..],
&LONG_VALUE1.as_bytes()[1..],
"Expected file content to match"
);

Expand Down Expand Up @@ -631,8 +660,8 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> {
}
}

let digest1 = DigestInfo::try_new(HASH1, VALUE1.len())?;
let digest2 = DigestInfo::try_new(HASH2, VALUE2.len())?;
let digest1 = DigestInfo::try_new(HASH1, LONG_VALUE1.len())?;
let digest2 = DigestInfo::try_new(HASH2, LONG_VALUE2.len())?;
let content_path = make_temp_path("content_path");
let temp_path = make_temp_path("temp_path");

Expand All @@ -652,7 +681,10 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> {
);

// Insert data into store.
store.update_oneshot(digest1, VALUE1.into()).await.unwrap();
store
.update_oneshot(digest1, LONG_VALUE1.into())
.await
.unwrap();

let mut reader = {
let (writer, reader) = make_buf_channel_pair();
Expand All @@ -668,7 +700,7 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> {
assert!(reader.peek().await.is_ok(), "Could not peek into reader");

// Insert new content. This will evict the old item.
store.update_oneshot(digest2, VALUE2.into()).await?;
store.update_oneshot(digest2, LONG_VALUE2.into()).await?;

// Ensure we let any background tasks finish.
tokio::task::yield_now().await;
Expand All @@ -687,7 +719,7 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> {
let data = read_file_contents(path.as_os_str()).await?;
assert_eq!(
&data[..],
VALUE1.as_bytes(),
LONG_VALUE1.as_bytes(),
"Expected file content to match"
);
}
Expand All @@ -702,7 +734,7 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> {
.await
.err_tip(|| "Error reading remaining bytes")?;

assert_eq!(&reader_data, VALUE1, "Expected file content to match");
assert_eq!(&reader_data, LONG_VALUE1, "Expected file content to match");

loop {
if DELETES_FINISHED.load(Ordering::Relaxed) == 1 {
Expand Down
16 changes: 11 additions & 5 deletions nativelink-util/src/buf_channel.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,18 +27,24 @@ use tracing::warn;

const ZERO_DATA: Bytes = Bytes::new();

/// Chunks buffered in flight per channel. Small enough that a fast producer
/// still applies backpressure promptly, but not so small (it was 2) that a
/// momentarily-slow consumer — e.g. a shard CAS mid-forward of a multi-GiB blob
/// — stalls the producer immediately. That instant stall stops the server
/// draining its inbound stream, the HTTP/2 flow-control window empties, and the
/// connection resets mid-stream (dropping large cacheable uploads). tokio's
/// bounded channel does not preallocate, so this capacity costs memory only
/// while data is actually queued (i.e. exactly during a downstream stall).
const CHANNEL_BUFFER_CHUNKS: usize = 32;

/// Create a channel pair that can be used to transport buffer objects around to
/// different components. This wrapper is used because the streams give some
/// utility like managing EOF in a more friendly way, ensure if no EOF is received
/// it will send an error to the receiver channel before shutting down and count
/// the number of bytes sent.
#[must_use]
pub fn make_buf_channel_pair() -> (DropCloserWriteHalf, DropCloserReadHalf) {
// We allow up to 2 items in the buffer at any given time. There is no major
// reason behind this magic number other than thinking it will be nice to give
// a little time for another thread to wake up and consume data if another
// thread is pumping large amounts of data into the channel.
let (tx, rx) = mpsc::channel(2);
let (tx, rx) = mpsc::channel(CHANNEL_BUFFER_CHUNKS);
let eof_sent = Arc::new(AtomicBool::new(false));
(
DropCloserWriteHalf {
Expand Down
Loading