perf: concurrent chunked upload and download of a single file over SFTP (#3393)

yazi-sftp/src/fs/file.rs

@@ -1,21 +1,28 @@
-use std::{io, pin::Pin, sync::Arc, task::{Context, Poll, ready}, time::Duration};
+use std::{io::{self, SeekFrom}, pin::Pin, sync::Arc, task::{Context, Poll, ready}, time::Duration};
 
-use tokio::{io::{AsyncRead, AsyncWrite, ReadBuf}, time::{Timeout, timeout}};
+use tokio::{io::{AsyncRead, AsyncSeek, AsyncWrite, ReadBuf}, time::{Timeout, timeout}};
 
-use crate::{Error, Operator, Packet, Receiver, Session, fs::Attrs};
+use crate::{Error, Operator, Packet, Receiver, Session, fs::Attrs, requests};
 
 pub struct File {
 	session: Arc<Session>,
-	handle:  String,
 
-	closed:   bool,
-	cursor:   u64,
+	handle: String,
+	cursor: u64,
+	closed: bool,
+
 	close_rx: Option<Timeout<Receiver>>,
 	read_rx:  Option<Receiver>,
+	seek_rx:  Option<SeekState>,
 	write_rx: Option<(Receiver, usize)>,
 	flush_rx: Option<Timeout<Receiver>>,
 }
 
+enum SeekState {
+	NonBlocking(u64),
+	Blocking(i64, Timeout<Receiver>),
+}
+
 impl Unpin for File {}
 
 impl Drop for File {
@@ -30,12 +37,14 @@ impl File {
 	pub(crate) fn new(session: &Arc<Session>, handle: impl Into<String>) -> Self {
 		Self {
 			session: session.clone(),
-			handle:  handle.into(),
 
-			closed:   false,
-			cursor:   0,
+			handle: handle.into(),
+			closed: false,
+			cursor: 0,
+
 			close_rx: None,
 			read_rx:  None,
+			seek_rx:  None,
 			write_rx: None,
 			flush_rx: None,
 		}
@@ -81,6 +90,78 @@ impl AsyncRead for File {
 	}
 }
 
+impl AsyncSeek for File {
+	fn start_seek(mut self: Pin<&mut Self>, position: io::SeekFrom) -> io::Result<()> {
+		if self.seek_rx.is_some() {
+			return Err(io::Error::other(
+				"other file operation is pending, call poll_complete before start_seek",
+			));
+		}
+
+		self.seek_rx = Some(match position {
+			SeekFrom::Start(n) => SeekState::NonBlocking(n),
+			SeekFrom::Current(n) => self
+				.cursor
+				.checked_add_signed(n)
+				.map(SeekState::NonBlocking)
+				.ok_or_else(|| io::Error::other("seeking to a negative or overflowed position"))?,
+			SeekFrom::End(n) => SeekState::Blocking(
+				n,
+				timeout(
+					Duration::from_secs(10),
+					self.session.send_sync(requests::Fstat::new(&self.handle))?,
+				),
+			),
+		});
+
+		Ok(())
+	}
+
+	fn poll_complete(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
+		let me = unsafe { self.get_unchecked_mut() };
+
+		let Some(state) = &mut me.seek_rx else {
+			return Poll::Ready(Ok(me.cursor));
+		};
+
+		fn imp(cx: &mut Context<'_>, state: &mut SeekState) -> Poll<io::Result<u64>> {
+			use Poll::Ready;
+			let (n, rx) = match state {
+				SeekState::NonBlocking(n) => return Ready(Ok(*n)),
+				SeekState::Blocking(n, rx) => (n, rx),
+			};
+
+			let Ok(result) = ready!(unsafe { Pin::new_unchecked(rx) }.poll(cx)) else {
+				return Ready(Err(Error::Timeout.into()));
+			};
+
+			let packet = match result {
+				Ok(Packet::Attrs(packet)) => packet,
+				Ok(_) => return Ready(Err(Error::Packet("not an Attrs").into())),
+				Err(e) => return Ready(Err(Error::from(e).into())),
+			};
+
+			let Some(size) = packet.attrs.size else {
+				return Ready(Err(io::Error::other("could not get file size for seeking from end")));
+			};
+
+			Ready(
+				size
+					.checked_add_signed(*n)
+					.ok_or_else(|| io::Error::other("seeking to a negative or overflowed position")),
+			)
+		}
+
+		let result = ready!(imp(cx, state));
+		if let Ok(n) = result {
+			me.cursor = n;
+		}
+
+		me.seek_rx = None;
+		Poll::Ready(result)
+	}
+}
+
 impl AsyncWrite for File {
 	fn poll_write(
 		self: Pin<&mut Self>,