use std::{
path::Path,
sync::{Arc, atomic::AtomicU64},
thread,
};
use anyhow::{Context, Result, bail};
use async_trait::async_trait;
use reqwest::Url;
use super::{DataReaderTrait, network_reader::NetworkReader, sftp_pool, sftp_utils};
use crate::{Blob, ByteRange};
pub struct DataReaderSftp {
connection: Arc<sftp_pool::Connection>,
file_id: u64,
size: u64,
name: String,
max_request_bytes: AtomicU64,
}
impl std::fmt::Debug for DataReaderSftp {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DataReaderSftp").field("name", &self.name).finish()
}
}
impl DataReaderSftp {
pub fn open(url: &Url, identity_file: Option<&Path>) -> Result<DataReaderSftp> {
let name = sftp_utils::display_name(url);
let connection = sftp_pool::acquire(url, identity_file)?;
let path = sftp_utils::remote_path(url);
let (file_id, size) = connection
.register(&path)
.with_context(|| format!("failed to open '{name}'"))?;
Ok(DataReaderSftp {
connection,
file_id,
size,
name,
max_request_bytes: AtomicU64::new(u64::MAX),
})
}
}
impl Drop for DataReaderSftp {
fn drop(&mut self) {
self.connection.unregister(self.file_id);
}
}
impl TryFrom<&Url> for DataReaderSftp {
type Error = anyhow::Error;
fn try_from(url: &Url) -> Result<DataReaderSftp> {
Self::open(url, None)
}
}
impl DataReaderSftp {
fn try_read_range_impl(&self, range: &ByteRange) -> Result<Blob> {
let policy = super::retry::policy();
let max_retries = policy.max_retries;
let total_attempts = max_retries + 1;
let name = &self.name;
let len = range.length;
let mut generation = self.connection.generation()?;
for attempt in 0..=max_retries {
let attempt_label = format!("attempt {}/{total_attempts}", attempt + 1);
if attempt > 0 {
let backoff = policy.backoff(attempt - 1);
log::warn!("SFTP read {range} from '{name}': retrying ({attempt_label}, waiting {backoff:?})");
thread::sleep(backoff);
if let Err(e) = self.connection.reconnect(generation) {
log::warn!("SFTP read {range} from '{name}': reconnect failed ({attempt_label}): {e}");
if attempt >= max_retries {
return Err(e).with_context(|| {
format!(
"could not read {range} ({len} bytes) from '{name}': reconnect failed — gave up after {total_attempts} attempts"
)
});
}
continue;
}
}
generation = self.connection.generation()?;
match self.connection.read_range(self.file_id, range) {
Ok(blob) => return Ok(blob),
Err(e) if attempt < max_retries => {
log::warn!("SFTP read {range} from '{name}': {e} ({attempt_label}), will retry");
}
Err(e) => {
return Err(e).with_context(|| {
format!(
"could not read {range} ({len} bytes) from '{name}' — gave up after {total_attempts} attempts"
)
});
}
}
}
bail!("SFTP read retry loop exited without returning — MAX_RETRIES invariant violated")
}
}
#[async_trait]
impl NetworkReader for DataReaderSftp {
async fn try_read_range(&self, range: &ByteRange) -> Result<Blob> {
self.try_read_range_impl(range)
}
fn max_request_bytes(&self) -> &AtomicU64 {
&self.max_request_bytes
}
}
#[async_trait]
impl DataReaderTrait for DataReaderSftp {
async fn read_range(&self, range: &ByteRange) -> Result<Blob> {
self.network_read_range(range).await
}
async fn read_all(&self) -> Result<Blob> {
self.read_range(&ByteRange::new(0, self.size)).await
}
fn name(&self) -> &str {
&self.name
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_open_unreachable_host() {
let url = Url::parse("sftp://192.0.2.1:22222/path/file.versatiles").unwrap();
let result = DataReaderSftp::open(&url, None);
assert!(result.is_err());
}
#[test]
fn test_try_from_unreachable_host() {
let url = Url::parse("sftp://192.0.2.1:22222/path/file.versatiles").unwrap();
let result = DataReaderSftp::try_from(&url);
assert!(result.is_err());
}
#[test]
fn test_open_with_identity_file_unreachable() {
let url = Url::parse("sftp://192.0.2.1:22222/path/file.versatiles").unwrap();
let result = DataReaderSftp::open(&url, Some(Path::new("/nonexistent/key")));
assert!(result.is_err());
}
#[cfg(all(feature = "ssh2", unix))]
mod sftp_server_tests {
use super::*;
use crate::{ByteRange, io::test_sftp_server::TestSftpServer};
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn read_all_returns_correct_bytes() {
let server = TestSftpServer::start().await;
let data: Vec<u8> = (0u8..100).collect();
server.write_file("/test.bin", &data).await;
let url = server.url("/test.bin");
let reader = tokio::task::spawn_blocking(move || DataReaderSftp::open(&url, None))
.await
.unwrap()
.unwrap();
let result = reader.read_all().await.unwrap();
assert_eq!(result.as_slice(), data.as_slice());
}
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn read_range_returns_slice() {
let server = TestSftpServer::start().await;
let data: Vec<u8> = (0u8..100).collect();
server.write_file("/test.bin", &data).await;
let url = server.url("/test.bin");
let reader = tokio::task::spawn_blocking(move || DataReaderSftp::open(&url, None))
.await
.unwrap()
.unwrap();
let result = reader.read_range(&ByteRange::new(20, 30)).await.unwrap();
assert_eq!(result.as_slice(), &data[20..50]);
}
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn read_range_at_eof() {
let server = TestSftpServer::start().await;
server.write_file("/test.bin", b"hello").await;
let url = server.url("/test.bin");
let reader = tokio::task::spawn_blocking(move || DataReaderSftp::open(&url, None))
.await
.unwrap()
.unwrap();
let result = reader.read_range(&ByteRange::new(5, 0)).await.unwrap();
assert!(result.is_empty());
}
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn read_retry_after_disconnect() {
let server = TestSftpServer::start().await;
let data: Vec<u8> = (0u8..50).collect();
server.write_file("/test.bin", &data).await;
let url = server.url("/test.bin");
let reader = tokio::task::spawn_blocking(move || DataReaderSftp::open(&url, None))
.await
.unwrap()
.unwrap();
server.schedule_disconnect();
let result = reader.read_range(&ByteRange::new(0, 50)).await.unwrap();
assert_eq!(result.as_slice(), data.as_slice());
}
#[tokio::test(flavor = "current_thread")]
#[serial_test::serial]
async fn open_nonexistent_file_errors() {
let server = TestSftpServer::start().await;
let url = server.url("/nonexistent.bin");
let result = tokio::task::spawn_blocking(move || DataReaderSftp::open(&url, None))
.await
.unwrap();
assert!(result.is_err());
}
}
}