Fix bug in buf_channel::consume() where exact size doesn't receive eof

nativelink-store/tests/filesystem_store_test.rs

@@ -228,8 +228,8 @@ mod filesystem_store_tests {
 
     const HASH1: &str = "0123456789abcdef000000000000000000010000000000000123456789abcdef";
     const HASH2: &str = "0123456789abcdef000000000000000000020000000000000123456789abcdef";
-    const VALUE1: &str = "1234";
-    const VALUE2: &str = "4321";
+    const VALUE1: &str = "0123456789";
+    const VALUE2: &str = "9876543210";
 
     #[tokio::test]
     async fn valid_results_after_shutdown_test() -> Result<(), Error> {
@@ -814,8 +814,9 @@ mod filesystem_store_tests {
 
     #[tokio::test]
     async fn eviction_on_insert_calls_unref_once() -> Result<(), Error> {
+        const SMALL_VALUE: &str = "01";
         const BIG_VALUE: &str = "0123";
-        let small_digest = DigestInfo::try_new(HASH1, VALUE1.len())?;
+        let small_digest = DigestInfo::try_new(HASH1, SMALL_VALUE.len())?;
         let big_digest = DigestInfo::try_new(HASH1, BIG_VALUE.len())?;
 
         static UNREFED_DIGESTS: Lazy<Mutex<Vec<DigestInfo>>> = Lazy::new(|| Mutex::new(Vec::new()));
@@ -851,7 +852,7 @@ mod filesystem_store_tests {
         // Insert data into store.
         store
             .as_ref()
-            .update_oneshot(small_digest, VALUE1.into())
+            .update_oneshot(small_digest, SMALL_VALUE.into())
             .await?;
         store
             .as_ref()

nativelink-util/src/buf_channel.rs

@@ -218,7 +218,7 @@ impl DropCloserReadHalf {
                 if !self.eof_sent.load(Ordering::Acquire) {
                     return Err(make_err!(
                         Code::Internal,
-                        "EOF received before sending EOF; sender was probably dropped"
+                        "Sender dropped before sending EOF"
                     ));
                 }
                 self.maybe_populate_recent_data(&ZERO_DATA);
@@ -323,14 +323,13 @@ impl DropCloserReadHalf {
             if chunk.len() > size {
                 let remaining = chunk.split_off(size);
                 self.queued_data.push_front(Ok(remaining));
-            }
-            if chunk.len() == size {
+                // No need to read EOF if we are a partial chunk.
                 return Ok(chunk);
             }
-            // If we are a partial chunk and our next chunk is EOF, we are done.
+            // Try to read our EOF to ensure our sender did not error out.
             match self.peek().await {
                 Ok(peeked_chunk) => {
-                    if peeked_chunk.is_empty() {
+                    if peeked_chunk.is_empty() || chunk.len() == size {
                         return Ok(chunk);
                     }
                 }

nativelink-util/tests/buf_channel_test.rs

@@ -12,7 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::task::Poll;
+
 use bytes::{Bytes, BytesMut};
+use futures::poll;
 use nativelink_error::{make_err, Code, Error, ResultExt};
 use nativelink_util::buf_channel::make_buf_channel_pair;
 use tokio::try_join;
@@ -67,7 +70,7 @@ mod buf_channel_tests {
     }
 
     #[tokio::test]
-    async fn consume_test() -> Result<(), Error> {
+    async fn consume_all_test() -> Result<(), Error> {
         let (mut tx, mut rx) = make_buf_channel_pair();
         let tx_fut = async move {
             tx.send(DATA1.into()).await?;
@@ -91,7 +94,7 @@ mod buf_channel_tests {
     /// Test to ensure data is optimized so that the exact same pointer is received
     /// when calling `collect_all_with_size_hint` when a copy is not needed.
     #[tokio::test]
-    async fn consume_is_optimized_test() -> Result<(), Error> {
+    async fn consume_all_is_optimized_test() -> Result<(), Error> {
         let (mut tx, mut rx) = make_buf_channel_pair();
         let sent_data = Bytes::from(DATA1);
         let send_data_ptr = sent_data.as_ptr();
@@ -112,7 +115,7 @@ mod buf_channel_tests {
     }
 
     #[tokio::test]
-    async fn take_test() -> Result<(), Error> {
+    async fn consume_some_test() -> Result<(), Error> {
         let (mut tx, mut rx) = make_buf_channel_pair();
         let tx_fut = async move {
             tx.send(DATA1.into()).await?;
@@ -139,7 +142,7 @@ mod buf_channel_tests {
     /// we don't need to concat the data together and instead return a view to
     /// the original data instead of making a copy.
     #[tokio::test]
-    async fn take_optimized_test() -> Result<(), Error> {
+    async fn consume_some_optimized_test() -> Result<(), Error> {
         let (mut tx, mut rx) = make_buf_channel_pair();
         let first_chunk = Bytes::from(DATA1);
         let first_chunk_ptr = first_chunk.as_ptr();
@@ -159,6 +162,23 @@ mod buf_channel_tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn consume_some_reads_eof() -> Result<(), Error> {
+        let (mut tx, mut rx) = make_buf_channel_pair();
+        tx.send(DATA1.into()).await?;
+
+        let consume_fut = rx.consume(Some(DATA1.len()));
+        tokio::pin!(consume_fut);
+        assert_eq!(
+            poll!(&mut consume_fut),
+            Poll::Pending,
+            "Consume should not have completed yet"
+        );
+        tx.send_eof()?;
+        assert_eq!(consume_fut.await?, Bytes::from(DATA1));
+        Ok(())
+    }
+
     #[tokio::test]
     async fn simple_stream_test() -> Result<(), Error> {
         use futures::StreamExt;
@@ -245,7 +265,7 @@ mod buf_channel_tests {
                 rx.recv().await,
                 Err(make_err!(
                     Code::Internal,
-                    "EOF received before sending EOF; sender was probably dropped"
+                    "Sender dropped before sending EOF"
                 ))
             );
             Result::<(), Error>::Ok(())