io: add AsyncSeek trait (#1924)

Co-authored-by: Taiki Endo <te316e89@gmail.com>

1 files changed, 77 insertions, 58 deletions

diff --git a/tokio/src/fs/file.rs b/tokio/src/fs/file.rs
index bb1084a1..20139ec1 100644
--- a/tokio/src/fs/file.rs
+++ b/tokio/src/fs/file.rs
@@ -5,7 +5,7 @@
 use self::State::*;
 use crate::fs::{asyncify, sys};
 use crate::io::blocking::Buf;
-use crate::io::{AsyncRead, AsyncWrite};
+use crate::io::{AsyncRead, AsyncSeek, AsyncWrite};
 
 use std::fmt;
 use std::fs::{Metadata, Permissions};
@@ -176,63 +176,6 @@ impl File {
         }
     }
 
-    /// Seek to an offset, in bytes, in a stream.
-    ///
-    /// # Examples
-    ///
-    /// ```no_run
-    /// use tokio::fs::File;
-    /// use tokio::prelude::*;
-    ///
-    /// use std::io::SeekFrom;
-    ///
-    /// # async fn dox() -> std::io::Result<()> {
-    /// let mut file = File::open("foo.txt").await?;
-    /// file.seek(SeekFrom::Start(6)).await?;
-    ///
-    /// let mut contents = vec![0u8; 10];
-    /// file.read_exact(&mut contents).await?;
-    /// # Ok(())
-    /// # }
-    /// ```
-    pub async fn seek(&mut self, mut pos: SeekFrom) -> io::Result<u64> {
-        self.complete_inflight().await;
-
-        let mut buf = match self.state {
-            Idle(ref mut buf_cell) => buf_cell.take().unwrap(),
-            _ => unreachable!(),
-        };
-
-        // Factor in any unread data from the buf
-        if !buf.is_empty() {
-            let n = buf.discard_read();
-
-            if let SeekFrom::Current(ref mut offset) = pos {
-                *offset += n;
-            }
-        }
-
-        let std = self.std.clone();
-
-        // Start the operation
-        self.state = Busy(sys::run(move || {
-            let res = (&*std).seek(pos);
-            (Operation::Seek(res), buf)
-        }));
-
-        let (op, buf) = match self.state {
-            Idle(_) => unreachable!(),
-            Busy(ref mut rx) => rx.await.unwrap(),
-        };
-
-        self.state = Idle(Some(buf));
-
-        match op {
-            Operation::Seek(res) => res,
-            _ => unreachable!(),
-        }
-    }
-
     /// Attempts to sync all OS-internal metadata to disk.
     ///
     /// This function will attempt to ensure that all in-core data reaches the
@@ -548,6 +491,82 @@ impl AsyncRead for File {
     }
 }
 
+impl AsyncSeek for File {
+    fn start_seek(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        mut pos: SeekFrom,
+    ) -> Poll<io::Result<()>> {
+        if let Some(e) = self.last_write_err.take() {
+            return Ready(Err(e.into()));
+        }
+
+        loop {
+            match self.state {
+                Idle(ref mut buf_cell) => {
+                    let mut buf = buf_cell.take().unwrap();
+
+                    // Factor in any unread data from the buf
+                    if !buf.is_empty() {
+                        let n = buf.discard_read();
+
+                        if let SeekFrom::Current(ref mut offset) = pos {
+                            *offset += n;
+                        }
+                    }
+
+                    let std = self.std.clone();
+
+                    self.state = Busy(sys::run(move || {
+                        let res = (&*std).seek(pos);
+                        (Operation::Seek(res), buf)
+                    }));
+
+                    return Ready(Ok(()));
+                }
+                Busy(ref mut rx) => {
+                    let (op, buf) = ready!(Pin::new(rx).poll(cx))?;
+                    self.state = Idle(Some(buf));
+
+                    match op {
+                        Operation::Read(_) => {}
+                        Operation::Write(Err(e)) => {
+                            self.last_write_err = Some(e.kind());
+                        }
+                        Operation::Write(_) => {}
+                        Operation::Seek(_) => {}
+                    }
+                }
+            }
+        }
+    }
+
+    fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
+        if let Some(e) = self.last_write_err.take() {
+            return Ready(Err(e.into()));
+        }
+
+        loop {
+            match self.state {
+                Idle(_) => panic!("must call start_seek before calling poll_complete"),
+                Busy(ref mut rx) => {
+                    let (op, buf) = ready!(Pin::new(rx).poll(cx))?;
+                    self.state = Idle(Some(buf));
+
+                    match op {
+                        Operation::Read(_) => {}
+                        Operation::Write(Err(e)) => {
+                            self.last_write_err = Some(e.kind());
+                        }
+                        Operation::Write(_) => {}
+                        Operation::Seek(res) => return Ready(res),
+                    }
+                }
+            }
+        }
+    }
+}
+
 impl AsyncWrite for File {
     fn poll_write(
         mut self: Pin<&mut Self>,