Skip to content

Commit

Permalink
Fix rare case where eof was sent on buf_channel when retry happens
Browse files Browse the repository at this point in the history
In a rare case if an EOF is sent, then a retry is triggered by
contract an EOF needs to be sent again, but without this patch
we are not allowed to send and EOF after one is already sent.
  • Loading branch information
allada committed Aug 30, 2024
1 parent f4ae4cc commit 1c88ead
Showing 1 changed file with 9 additions and 5 deletions.
14 changes: 9 additions & 5 deletions nativelink-util/src/buf_channel.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ use futures::task::Context;
use futures::{Future, Stream, TryFutureExt};
use nativelink_error::{error_if, make_err, make_input_err, Code, Error, ResultExt};
use tokio::sync::mpsc;
use tracing::{event, Level};

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

Expand Down Expand Up @@ -151,12 +152,15 @@ impl DropCloserWriteHalf {
/// Sends an EOF (End of File) message to the receiver which will gracefully let the
/// stream know it has no more data. This will close the stream.
pub fn send_eof(&mut self) -> Result<(), Error> {
error_if!(
self.tx.is_none(),
"Tried to send an EOF when pipe is broken"
);
// Flag that we have sent the EOF.
self.eof_sent.store(true, Ordering::Release);
let eof_was_sent = self.eof_sent.swap(true, Ordering::Release);
if eof_was_sent {
event!(
Level::WARN,
"Stream already closed when eof already was sent. This is often ok for retry was triggered, but should not happen on happy path."
);
return Ok(());
}

// Now close our stream.
self.tx = None;
Expand Down

0 comments on commit 1c88ead

Please sign in to comment.