Auto merge of #79780 - camelid:use-summary_opts, r=GuillaumeGomez

[rust.git] / library / std / src / sys / unix / kernel_copy.rs

diff --git a/library/std/src/sys/unix/kernel_copy.rs b/library/std/src/sys/unix/kernel_copy.rs
index 1dc16ef099367c8835f503943bc1e869d4b21a5d..5bfac80315348b959a5671805ed0fcc115514459 100644 (file)
--- a/library/std/src/sys/unix/kernel_copy.rs
+++ b/library/std/src/sys/unix/kernel_copy.rs
@@ -56,6 +56,7 @@
 use crate::net::TcpStream;
 use crate::os::unix::fs::FileTypeExt;
 use crate::os::unix::io::{AsRawFd, FromRawFd, RawFd};
+use crate::os::unix::net::UnixStream;
 use crate::process::{ChildStderr, ChildStdin, ChildStdout};
 use crate::ptr;
 use crate::sync::atomic::{AtomicBool, Ordering};
@@ -166,10 +167,11 @@ fn copy(self) -> Result<u64> {
 
             if input_meta.copy_file_range_candidate() && output_meta.copy_file_range_candidate() {
                 let result = copy_regular_files(readfd, writefd, max_write);
+                result.update_take(reader);
 
                 match result {
-                    CopyResult::Ended(Ok(bytes_copied)) => return Ok(bytes_copied + written),
-                    CopyResult::Ended(err) => return err,
+                    CopyResult::Ended(bytes_copied) => return Ok(bytes_copied + written),
+                    CopyResult::Error(e, _) => return Err(e),
                     CopyResult::Fallback(bytes) => written += bytes,
                 }
             }
@@ -181,20 +183,22 @@ fn copy(self) -> Result<u64> {
             // fall back to the generic copy loop.
             if input_meta.potential_sendfile_source() {
                 let result = sendfile_splice(SpliceMode::Sendfile, readfd, writefd, max_write);
+                result.update_take(reader);
 
                 match result {
-                    CopyResult::Ended(Ok(bytes_copied)) => return Ok(bytes_copied + written),
-                    CopyResult::Ended(err) => return err,
+                    CopyResult::Ended(bytes_copied) => return Ok(bytes_copied + written),
+                    CopyResult::Error(e, _) => return Err(e),
                     CopyResult::Fallback(bytes) => written += bytes,
                 }
             }
 
             if input_meta.maybe_fifo() || output_meta.maybe_fifo() {
                 let result = sendfile_splice(SpliceMode::Splice, readfd, writefd, max_write);
+                result.update_take(reader);
 
                 match result {
-                    CopyResult::Ended(Ok(bytes_copied)) => return Ok(bytes_copied + written),
-                    CopyResult::Ended(err) => return err,
+                    CopyResult::Ended(bytes_copied) => return Ok(bytes_copied + written),
+                    CopyResult::Error(e, _) => return Err(e),
                     CopyResult::Fallback(0) => { /* use the fallback below */ }
                     CopyResult::Fallback(_) => {
                         unreachable!("splice should not return > 0 bytes on the fallback path")
@@ -224,6 +228,9 @@ fn drain_to<W: Write>(&mut self, _writer: &mut W, _limit: u64) -> Result<u64> {
         Ok(0)
     }
 
+    /// Updates `Take` wrappers to remove the number of bytes copied.
+    fn taken(&mut self, _bytes: u64) {}
+
     /// The minimum of the limit of all `Take<_>` wrappers, `u64::MAX` otherwise.
     /// This method does not account for data `BufReader` buffers and would underreport
     /// the limit of a `Take<BufReader<Take<_>>>` type. Thus its result is only valid
@@ -250,6 +257,10 @@ fn drain_to<W: Write>(&mut self, writer: &mut W, limit: u64) -> Result<u64> {
         (**self).drain_to(writer, limit)
     }
 
+    fn taken(&mut self, bytes: u64) {
+        (**self).taken(bytes);
+    }
+
     fn min_limit(&self) -> u64 {
         (**self).min_limit()
     }
@@ -320,6 +331,34 @@ fn properties(&self) -> CopyParams {
     }
 }
 
+impl CopyRead for UnixStream {
+    fn properties(&self) -> CopyParams {
+        // avoid the stat syscall since we can be fairly sure it's a socket
+        CopyParams(FdMeta::Socket, Some(self.as_raw_fd()))
+    }
+}
+
+impl CopyRead for &UnixStream {
+    fn properties(&self) -> CopyParams {
+        // avoid the stat syscall since we can be fairly sure it's a socket
+        CopyParams(FdMeta::Socket, Some(self.as_raw_fd()))
+    }
+}
+
+impl CopyWrite for UnixStream {
+    fn properties(&self) -> CopyParams {
+        // avoid the stat syscall since we can be fairly sure it's a socket
+        CopyParams(FdMeta::Socket, Some(self.as_raw_fd()))
+    }
+}
+
+impl CopyWrite for &UnixStream {
+    fn properties(&self) -> CopyParams {
+        // avoid the stat syscall since we can be fairly sure it's a socket
+        CopyParams(FdMeta::Socket, Some(self.as_raw_fd()))
+    }
+}
+
 impl CopyWrite for ChildStdin {
     fn properties(&self) -> CopyParams {
         CopyParams(FdMeta::Pipe, Some(self.as_raw_fd()))
@@ -378,6 +417,11 @@ fn drain_to<W: Write>(&mut self, writer: &mut W, outer_limit: u64) -> Result<u64
         Ok(bytes_drained)
     }
 
+    fn taken(&mut self, bytes: u64) {
+        self.set_limit(self.limit() - bytes);
+        self.get_mut().taken(bytes);
+    }
+
     fn min_limit(&self) -> u64 {
         min(Take::limit(self), self.get_ref().min_limit())
     }
@@ -403,6 +447,10 @@ fn drain_to<W: Write>(&mut self, writer: &mut W, outer_limit: u64) -> Result<u64
         Ok(bytes as u64 + inner_bytes)
     }
 
+    fn taken(&mut self, bytes: u64) {
+        self.get_mut().taken(bytes);
+    }
+
     fn min_limit(&self) -> u64 {
         self.get_ref().min_limit()
     }
@@ -428,10 +476,21 @@ fn fd_to_meta<T: AsRawFd>(fd: &T) -> FdMeta {
 }
 
 pub(super) enum CopyResult {
-    Ended(Result<u64>),
+    Ended(u64),
+    Error(Error, u64),
     Fallback(u64),
 }
 
+impl CopyResult {
+    fn update_take(&self, reader: &mut impl CopyRead) {
+        match *self {
+            CopyResult::Fallback(bytes)
+            | CopyResult::Ended(bytes)
+            | CopyResult::Error(_, bytes) => reader.taken(bytes),
+        }
+    }
+}
+
 /// linux-specific implementation that will attempt to use copy_file_range for copy offloading
 /// as the name says, it only works on regular files
 ///
@@ -498,7 +557,7 @@ fn copy_file_range(
                 // - copying from an overlay filesystem in docker. reported to occur on fedora 32.
                 return CopyResult::Fallback(0);
             }
-            Ok(0) => return CopyResult::Ended(Ok(written)), // reached EOF
+            Ok(0) => return CopyResult::Ended(written), // reached EOF
             Ok(ret) => written += ret as u64,
             Err(err) => {
                 return match err.raw_os_error() {
@@ -516,12 +575,12 @@ fn copy_file_range(
                         assert_eq!(written, 0);
                         CopyResult::Fallback(0)
                     }
-                    _ => CopyResult::Ended(Err(err)),
+                    _ => CopyResult::Error(err, written),
                 };
             }
         }
     }
-    CopyResult::Ended(Ok(written))
+    CopyResult::Ended(written)
 }
 
 #[derive(PartialEq)]
@@ -594,10 +653,10 @@ fn splice(
                     Some(os_err) if mode == SpliceMode::Sendfile && os_err == libc::EOVERFLOW => {
                         CopyResult::Fallback(written)
                     }
-                    _ => CopyResult::Ended(Err(err)),
+                    _ => CopyResult::Error(err, written),
                 };
             }
         }
     }
-    CopyResult::Ended(Ok(written))
+    CopyResult::Ended(written)
 }