library/std/src/io/buffered/bufwriter.rs

   1 use crate::fmt;
   2 use crate::io::{
   3     self, Error, ErrorKind, IntoInnerError, IoSlice, Seek, SeekFrom, Write, DEFAULT_BUF_SIZE,
   4 };
   5
   6 /// Wraps a writer and buffers its output.
   7 ///
   8 /// It can be excessively inefficient to work directly with something that
   9 /// implements [`Write`]. For example, every call to
  10 /// [`write`][`TcpStream::write`] on [`TcpStream`] results in a system call. A
  11 /// `BufWriter<W>` keeps an in-memory buffer of data and writes it to an underlying
  12 /// writer in large, infrequent batches.
  13 ///
  14 /// `BufWriter<W>` can improve the speed of programs that make *small* and
  15 /// *repeated* write calls to the same file or network socket. It does not
  16 /// help when writing very large amounts at once, or writing just one or a few
  17 /// times. It also provides no advantage when writing to a destination that is
  18 /// in memory, like a [`Vec`]`<u8>`.
  19 ///
  20 /// It is critical to call [`flush`] before `BufWriter<W>` is dropped. Though
  21 /// dropping will attempt to flush the contents of the buffer, any errors
  22 /// that happen in the process of dropping will be ignored. Calling [`flush`]
  23 /// ensures that the buffer is empty and thus dropping will not even attempt
  24 /// file operations.
  25 ///
  26 /// # Examples
  27 ///
  28 /// Let's write the numbers one through ten to a [`TcpStream`]:
  29 ///
  30 /// ```no_run
  31 /// use std::io::prelude::*;
  32 /// use std::net::TcpStream;
  33 ///
  34 /// let mut stream = TcpStream::connect("127.0.0.1:34254").unwrap();
  35 ///
  36 /// for i in 0..10 {
  37 ///     stream.write(&[i+1]).unwrap();
  38 /// }
  39 /// ```
  40 ///
  41 /// Because we're not buffering, we write each one in turn, incurring the
  42 /// overhead of a system call per byte written. We can fix this with a
  43 /// `BufWriter<W>`:
  44 ///
  45 /// ```no_run
  46 /// use std::io::prelude::*;
  47 /// use std::io::BufWriter;
  48 /// use std::net::TcpStream;
  49 ///
  50 /// let mut stream = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
  51 ///
  52 /// for i in 0..10 {
  53 ///     stream.write(&[i+1]).unwrap();
  54 /// }
  55 /// stream.flush().unwrap();
  56 /// ```
  57 ///
  58 /// By wrapping the stream with a `BufWriter<W>`, these ten writes are all grouped
  59 /// together by the buffer and will all be written out in one system call when
  60 /// the `stream` is flushed.
  61 ///
  62 /// [`TcpStream::write`]: Write::write
  63 /// [`TcpStream`]: crate::net::TcpStream
  64 /// [`flush`]: Write::flush
  65 #[stable(feature = "rust1", since = "1.0.0")]
  66 pub struct BufWriter<W: Write> {
  67     inner: Option<W>,
  68     buf: Vec<u8>,
  69     // #30888: If the inner writer panics in a call to write, we don't want to
  70     // write the buffered data a second time in BufWriter's destructor. This
  71     // flag tells the Drop impl if it should skip the flush.
  72     panicked: bool,
  73 }
  74
  75 impl<W: Write> BufWriter<W> {
  76     /// Creates a new `BufWriter<W>` with a default buffer capacity. The default is currently 8 KB,
  77     /// but may change in the future.
  78     ///
  79     /// # Examples
  80     ///
  81     /// ```no_run
  82     /// use std::io::BufWriter;
  83     /// use std::net::TcpStream;
  84     ///
  85     /// let mut buffer = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
  86     /// ```
  87     #[stable(feature = "rust1", since = "1.0.0")]
  88     pub fn new(inner: W) -> BufWriter<W> {
  89         BufWriter::with_capacity(DEFAULT_BUF_SIZE, inner)
  90     }
  91
  92     /// Creates a new `BufWriter<W>` with the specified buffer capacity.
  93     ///
  94     /// # Examples
  95     ///
  96     /// Creating a buffer with a buffer of a hundred bytes.
  97     ///
  98     /// ```no_run
  99     /// use std::io::BufWriter;
 100     /// use std::net::TcpStream;
 101     ///
 102     /// let stream = TcpStream::connect("127.0.0.1:34254").unwrap();
 103     /// let mut buffer = BufWriter::with_capacity(100, stream);
 104     /// ```
 105     #[stable(feature = "rust1", since = "1.0.0")]
 106     pub fn with_capacity(capacity: usize, inner: W) -> BufWriter<W> {
 107         BufWriter { inner: Some(inner), buf: Vec::with_capacity(capacity), panicked: false }
 108     }
 109
 110     /// Send data in our local buffer into the inner writer, looping as
 111     /// necessary until either it's all been sent or an error occurs.
 112     ///
 113     /// Because all the data in the buffer has been reported to our owner as
 114     /// "successfully written" (by returning nonzero success values from
 115     /// `write`), any 0-length writes from `inner` must be reported as i/o
 116     /// errors from this method.
 117     pub(super) fn flush_buf(&mut self) -> io::Result<()> {
 118         /// Helper struct to ensure the buffer is updated after all the writes
 119         /// are complete. It tracks the number of written bytes and drains them
 120         /// all from the front of the buffer when dropped.
 121         struct BufGuard<'a> {
 122             buffer: &'a mut Vec<u8>,
 123             written: usize,
 124         }
 125
 126         impl<'a> BufGuard<'a> {
 127             fn new(buffer: &'a mut Vec<u8>) -> Self {
 128                 Self { buffer, written: 0 }
 129             }
 130
 131             /// The unwritten part of the buffer
 132             fn remaining(&self) -> &[u8] {
 133                 &self.buffer[self.written..]
 134             }
 135
 136             /// Flag some bytes as removed from the front of the buffer
 137             fn consume(&mut self, amt: usize) {
 138                 self.written += amt;
 139             }
 140
 141             /// true if all of the bytes have been written
 142             fn done(&self) -> bool {
 143                 self.written >= self.buffer.len()
 144             }
 145         }
 146
 147         impl Drop for BufGuard<'_> {
 148             fn drop(&mut self) {
 149                 if self.written > 0 {
 150                     self.buffer.drain(..self.written);
 151                 }
 152             }
 153         }
 154
 155         let mut guard = BufGuard::new(&mut self.buf);
 156         let inner = self.inner.as_mut().unwrap();
 157         while !guard.done() {
 158             self.panicked = true;
 159             let r = inner.write(guard.remaining());
 160             self.panicked = false;
 161
 162             match r {
 163                 Ok(0) => {
 164                     return Err(Error::new(
 165                         ErrorKind::WriteZero,
 166                         "failed to write the buffered data",
 167                     ));
 168                 }
 169                 Ok(n) => guard.consume(n),
 170                 Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
 171                 Err(e) => return Err(e),
 172             }
 173         }
 174         Ok(())
 175     }
 176
 177     /// Buffer some data without flushing it, regardless of the size of the
 178     /// data. Writes as much as possible without exceeding capacity. Returns
 179     /// the number of bytes written.
 180     pub(super) fn write_to_buf(&mut self, buf: &[u8]) -> usize {
 181         let available = self.buf.capacity() - self.buf.len();
 182         let amt_to_buffer = available.min(buf.len());
 183         self.buf.extend_from_slice(&buf[..amt_to_buffer]);
 184         amt_to_buffer
 185     }
 186
 187     /// Gets a reference to the underlying writer.
 188     ///
 189     /// # Examples
 190     ///
 191     /// ```no_run
 192     /// use std::io::BufWriter;
 193     /// use std::net::TcpStream;
 194     ///
 195     /// let mut buffer = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
 196     ///
 197     /// // we can use reference just like buffer
 198     /// let reference = buffer.get_ref();
 199     /// ```
 200     #[stable(feature = "rust1", since = "1.0.0")]
 201     pub fn get_ref(&self) -> &W {
 202         self.inner.as_ref().unwrap()
 203     }
 204
 205     /// Gets a mutable reference to the underlying writer.
 206     ///
 207     /// It is inadvisable to directly write to the underlying writer.
 208     ///
 209     /// # Examples
 210     ///
 211     /// ```no_run
 212     /// use std::io::BufWriter;
 213     /// use std::net::TcpStream;
 214     ///
 215     /// let mut buffer = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
 216     ///
 217     /// // we can use reference just like buffer
 218     /// let reference = buffer.get_mut();
 219     /// ```
 220     #[stable(feature = "rust1", since = "1.0.0")]
 221     pub fn get_mut(&mut self) -> &mut W {
 222         self.inner.as_mut().unwrap()
 223     }
 224
 225     /// Returns a reference to the internally buffered data.
 226     ///
 227     /// # Examples
 228     ///
 229     /// ```no_run
 230     /// use std::io::BufWriter;
 231     /// use std::net::TcpStream;
 232     ///
 233     /// let buf_writer = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
 234     ///
 235     /// // See how many bytes are currently buffered
 236     /// let bytes_buffered = buf_writer.buffer().len();
 237     /// ```
 238     #[stable(feature = "bufreader_buffer", since = "1.37.0")]
 239     pub fn buffer(&self) -> &[u8] {
 240         &self.buf
 241     }
 242
 243     /// Returns the number of bytes the internal buffer can hold without flushing.
 244     ///
 245     /// # Examples
 246     ///
 247     /// ```no_run
 248     /// use std::io::BufWriter;
 249     /// use std::net::TcpStream;
 250     ///
 251     /// let buf_writer = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
 252     ///
 253     /// // Check the capacity of the inner buffer
 254     /// let capacity = buf_writer.capacity();
 255     /// // Calculate how many bytes can be written without flushing
 256     /// let without_flush = capacity - buf_writer.buffer().len();
 257     /// ```
 258     #[stable(feature = "buffered_io_capacity", since = "1.46.0")]
 259     pub fn capacity(&self) -> usize {
 260         self.buf.capacity()
 261     }
 262
 263     /// Unwraps this `BufWriter<W>`, returning the underlying writer.
 264     ///
 265     /// The buffer is written out before returning the writer.
 266     ///
 267     /// # Errors
 268     ///
 269     /// An [`Err`] will be returned if an error occurs while flushing the buffer.
 270     ///
 271     /// # Examples
 272     ///
 273     /// ```no_run
 274     /// use std::io::BufWriter;
 275     /// use std::net::TcpStream;
 276     ///
 277     /// let mut buffer = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
 278     ///
 279     /// // unwrap the TcpStream and flush the buffer
 280     /// let stream = buffer.into_inner().unwrap();
 281     /// ```
 282     #[stable(feature = "rust1", since = "1.0.0")]
 283     pub fn into_inner(mut self) -> Result<W, IntoInnerError<BufWriter<W>>> {
 284         match self.flush_buf() {
 285             Err(e) => Err(IntoInnerError::new(self, e)),
 286             Ok(()) => Ok(self.inner.take().unwrap()),
 287         }
 288     }
 289 }
 290
 291 #[stable(feature = "rust1", since = "1.0.0")]
 292 impl<W: Write> Write for BufWriter<W> {
 293     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
 294         if self.buf.len() + buf.len() > self.buf.capacity() {
 295             self.flush_buf()?;
 296         }
 297         // FIXME: Why no len > capacity? Why not buffer len == capacity? #72919
 298         if buf.len() >= self.buf.capacity() {
 299             self.panicked = true;
 300             let r = self.get_mut().write(buf);
 301             self.panicked = false;
 302             r
 303         } else {
 304             self.buf.extend_from_slice(buf);
 305             Ok(buf.len())
 306         }
 307     }
 308
 309     fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
 310         // Normally, `write_all` just calls `write` in a loop. We can do better
 311         // by calling `self.get_mut().write_all()` directly, which avoids
 312         // round trips through the buffer in the event of a series of partial
 313         // writes in some circumstances.
 314         if self.buf.len() + buf.len() > self.buf.capacity() {
 315             self.flush_buf()?;
 316         }
 317         // FIXME: Why no len > capacity? Why not buffer len == capacity? #72919
 318         if buf.len() >= self.buf.capacity() {
 319             self.panicked = true;
 320             let r = self.get_mut().write_all(buf);
 321             self.panicked = false;
 322             r
 323         } else {
 324             self.buf.extend_from_slice(buf);
 325             Ok(())
 326         }
 327     }
 328
 329     fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
 330         let total_len = bufs.iter().map(|b| b.len()).sum::<usize>();
 331         if self.buf.len() + total_len > self.buf.capacity() {
 332             self.flush_buf()?;
 333         }
 334         // FIXME: Why no len > capacity? Why not buffer len == capacity? #72919
 335         if total_len >= self.buf.capacity() {
 336             self.panicked = true;
 337             let r = self.get_mut().write_vectored(bufs);
 338             self.panicked = false;
 339             r
 340         } else {
 341             bufs.iter().for_each(|b| self.buf.extend_from_slice(b));
 342             Ok(total_len)
 343         }
 344     }
 345
 346     fn is_write_vectored(&self) -> bool {
 347         self.get_ref().is_write_vectored()
 348     }
 349
 350     fn flush(&mut self) -> io::Result<()> {
 351         self.flush_buf().and_then(|()| self.get_mut().flush())
 352     }
 353 }
 354
 355 #[stable(feature = "rust1", since = "1.0.0")]
 356 impl<W: Write> fmt::Debug for BufWriter<W>
 357 where
 358     W: fmt::Debug,
 359 {
 360     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
 361         fmt.debug_struct("BufWriter")
 362             .field("writer", &self.inner.as_ref().unwrap())
 363             .field("buffer", &format_args!("{}/{}", self.buf.len(), self.buf.capacity()))
 364             .finish()
 365     }
 366 }
 367
 368 #[stable(feature = "rust1", since = "1.0.0")]
 369 impl<W: Write + Seek> Seek for BufWriter<W> {
 370     /// Seek to the offset, in bytes, in the underlying writer.
 371     ///
 372     /// Seeking always writes out the internal buffer before seeking.
 373     fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
 374         self.flush_buf()?;
 375         self.get_mut().seek(pos)
 376     }
 377 }
 378
 379 #[stable(feature = "rust1", since = "1.0.0")]
 380 impl<W: Write> Drop for BufWriter<W> {
 381     fn drop(&mut self) {
 382         if self.inner.is_some() && !self.panicked {
 383             // dtors should not panic, so we ignore a failed flush
 384             let _r = self.flush_buf();
 385         }
 386     }
 387 }